def substack()

in multiset_codec/codecs.py [0:0]

• 20 lines of code
• 1 McCabe index (conditional complexity)

def substack(codec: Codec, view_fun: ViewFunc) -> Codec:
    '''
    Apply a codec on a subset of a ans_state head.
    view_fun should be a function: head -> subhead, for example
    view_fun = lambda head: head[0]
    to run the codec on only the first element of the head
    '''
    def encode(ans_state, symbol, *context):
        head, tail = ans_state
        subhead, update = cs.util.view_update(head, view_fun)
        (subhead, tail), *context = \
                codec.encode((subhead, tail), symbol, *context)
        return ((update(subhead), tail), *context)

    def decode(ans_state, *context):
        head, tail = ans_state
        subhead, update = cs.util.view_update(head, view_fun)
        (subhead, tail), symbol, *context = \
                codec.decode((subhead, tail), *context)
        return ((update(subhead), tail), symbol, *context)

    return Codec(encode, decode)