Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Minor, Opinionated Cleanup for atproto/packages/atproto_firehose/models.py #511

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Minor, Opinionated Cleanup for atproto/packages/atproto_firehose/models.py #511

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0fbb979
FrameType Enum now checks for equality against header_op, neagting th…
PyWoody Dec 31, 2024
fa850b4
Clarified that Frame.frome_bytes raises 'atproto.exceptions.FirehoseD…
PyWoody Dec 31, 2024
4ce3df3
Simplified the logic for decoding parts in 'Frame.from_bytes'.
PyWoody Dec 31, 2024
9f4d135
Equality checks for Enum require checking against value
PyWoody Dec 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 10 additions & 22 deletions packages/atproto_firehose/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,6 @@ class FrameType(Enum):
MESSAGE = 1
ERROR = -1

@classmethod
def has_value(cls, value: int) -> bool:
return value in cls._value2member_map_


@dataclass
class MessageFrameHeader:
Expand Down Expand Up @@ -50,13 +46,12 @@ class ErrorFrameBody:
def parse_frame_header(raw_header: dict) -> FrameHeader:
try:
header_op = int(raw_header.get('op', 0))
if not FrameType.has_value(header_op):
raise FirehoseDecodingError('Invalid frame type')

frame_type = FrameType(header_op)
if frame_type is FrameType.MESSAGE:
if header_op == FrameType.MESSAGE.value:
return get_or_create(raw_header, MessageFrameHeader)
return get_or_create(raw_header, ErrorFrameHeader)
elif header_op == FrameType.ERROR.value:
return get_or_create(raw_header, ErrorFrameHeader)
else:
raise FirehoseDecodingError('Invalid frame type')
except (ValueError, AtProtocolError) as e:
raise FirehoseDecodingError('Invalid frame header') from e

Expand Down Expand Up @@ -106,23 +101,16 @@ def from_bytes(data: Union[bytes, bytearray]) -> Union['MessageFrame', 'ErrorFra
:obj:`atproto.firehose_models.MessageFrame` or :obj:`atproto.firehose_models.ErrorFrame`

Raises:
:class:`atproto.exceptions.FirehoseError`: Invalid data frame.
:class:`atproto.exceptions.FirehoseDecodingError`: Invalid data frame.
"""
decoded_parts = decode_dag_multi(data)
if not decoded_parts:
raise FirehoseDecodingError('Invalid frame without CBOR data')
if len(decoded_parts) > 2:
raise FirehoseDecodingError('Too many CBOR data parts in the frame')
if not len(decoded_parts):
raise FirehoseDecodingError('Invalid frame without CBOR data')

raw_header = decoded_parts[0]

raw_body = None
if len(decoded_parts) > 1:
raw_body = decoded_parts[1]

if raw_body is None:
if len(decoded_parts) < 2:
raise FirehoseDecodingError('Frame body not found')

raw_header, raw_body = decoded_parts
header = parse_frame_header(raw_header)
return parse_frame(header, raw_body)

Expand Down