Guidance for using If based on total length of data

[maihde]

I have the data that comes in two different lengths, lets say 1 and 2 bytes for simplicity where the longer one uses the same format as the shorter, but includes a prefix. A straight forward approach is to write something like this:

Short = BitStruct(
    "fieldA" / c.BitsInteger(4),
    "fieldB" / c.BitsInteger(4),
)

Long = BitStruct(
    "prefixA", / c.BitsInteger(2),
    "prefixB", / c.BitsInteger(4),
    "prefixC", / c.BitsInteger(2),
    "fieldA" / c.BitsInteger(4),
    "fieldB" / c.BitsInteger(4),
)

if len(data) == 1:
    return Short.parse(data)
else:
    return Long.parse(data)

But if felt like there was a more elegant solution

Prefix = BitStruct(
    "prefixA", / c.BitsInteger(2),
    "prefixB", / c.BitsInteger(4),
    "prefixC", / c.BitsInteger(2),
)

Message = BitStruct(
    "prefix" / If(lambda this: ??, Bytewise(Prefix))
    "fieldA" / c.BitsInteger(4),
    "fieldB" / c.BitsInteger(4),
)

return Message.parse(data)

I couldn't find a way to get the overall data length from the context to satisfy the If so I hacked something like this:

class DataLength(c.Tunnel):
    def _decode(self, data, context, path):
        context._datalen = len(data)
        return data
    
    def _encode(self, data, context, path):
        return data
        
Prefix = BitStruct(
    "prefixA", / c.BitsInteger(2),
    "prefixB", / c.BitsInteger(4),
    "prefixC", / c.BitsInteger(2),
)

Message = DataLength(
    BitStruct(
        "prefix" / If(lambda this: this._._datalen == 2, Bytewise(Prefix))
        "fieldA" / c.BitsInteger(4),
        "fieldB" / c.BitsInteger(4),
    )
)

return Message.parse(data)

It works, but feels like I'm missing something that should let me accomplish this same outcome without the custom Tunnel. I'd appreciate any guidance on the best way to proceed.

[arekbulski]

How about something like this, typing off my head:

Message = Struct(
    "data" / GreedyBytes,
    "parsed" / Switch(len_(this.data),
        1:Short,
        2:Long,
    ),
)