scott-griffiths / bitstring

A Python module to help you manage your bits
https://bitstring.readthedocs.io/en/stable/index.html
MIT License
404 stars 68 forks source link

Make typing of BitStream.read less ambigous #262

Closed prauscher closed 1 year ago

prauscher commented 1 year ago

Currently, the signature of BitStream.read is given as

def read(self, fmt: Union[int, str]) -> Union[int, float, str, Bits, bool, bytes, None]:

This leads to problems when using:

# error: Unsupported operand types for + ("int" and "str")
# error: Unsupported operand types for + ("int" and "bytes")
# error: Unsupported operand types for + ("int" and "None")
# note: Right operand is of type "Union[float, str, Bits, bytes, None]"
print(4 + bitstream.read("uint:4"))

# error: Item "float" of "Union[float, str, Bits, bytes, None]" has no attribute "uint"
# error: Item "str" of "Union[float, str, Bits, bytes, None]" has no attribute "uint"
# error: Item "bytes" of "Union[float, str, Bits, bytes, None]" has no attribute "uint"
# error: Item "None" of "Union[float, str, Bits, bytes, None]" has no attribute "uint"
print(4 + bitstream.read(4).uint)

Which is obviously invalid, as both calls will work with integers. One solution would be to use @overload

e.g.

@overload
def read(self, fmt: int) -> Bits:
    ...

def read(self, fmt: Union[int, str]) -> Union[int, float, str, Bits, bool, bytes, None]:
scott-griffiths commented 1 year ago

Hi. Yes you're right that there are probably a few methods that could use some more explicit overloading hints. Thanks for pointing this out.

scott-griffiths commented 1 year ago

After a bit of trial and error, the best I could do was to use an Any to sooth the mypy gods:

    @overload
    def read(self, fmt: int) -> Bits: ...
    @overload
    def read(self, fmt: str) -> Any: ...

    def read(self, fmt: Union[int, str]) -> Union[int, float, str, Bits, bool, bytes, None]:

I'm fine with that as there's no way to statically work out the return type if its reading a format str.

scott-griffiths commented 1 year ago

Done in the 4.1 beta now available.