soia-client 1.1.4__py3-none-any.whl → 1.1.6__py3-none-any.whl

soia/_impl/structs.py

@@ -1,8 +1,10 @@
 import copy
 from collections.abc import Callable, Sequence
 from dataclasses import FrozenInstanceError, dataclass
-from typing import Any, Final, Union, cast
+import itertools
+from typing import Any, Final, Generic, Union, cast
 
+from soia import _spec, reflection
 from soia._impl.function_maker import (
     BodyBuilder,
     Expr,
@@ -16,12 +18,11 @@ from soia._impl.function_maker import (
 )
 from soia._impl.keep import KEEP
 from soia._impl.repr import repr_impl
-from soia._impl.type_adapter import TypeAdapter
-
-from soia import _spec, reflection
+from soia._impl.type_adapter import T, ByteStream, TypeAdapter
+from soia._impl.binary import decode_int64, decode_unused, encode_length_prefix
 
 
-class StructAdapter(TypeAdapter):
+class StructAdapter(Generic[T], TypeAdapter[T]):
     __slots__ = (
         "spec",
         "record_hash",
@@ -44,6 +45,8 @@ class StructAdapter(TypeAdapter):
     # 0: has not started; 1: in progress; 2: done
     finalization_state: int
     fields: tuple["_Field", ...]
+    num_slots_incl_removed: int
+    slot_to_field: list["_Field | None"]
 
     def __init__(self, spec: _spec.Struct):
         self.finalization_state = 0
@@ -83,6 +86,17 @@ class StructAdapter(TypeAdapter):
         # Reason why it's faster: we don't need to call object.__setattr__().
         self.simple_class = _make_dataclass(slots)
 
+        def forward_encode(value: Any, buffer: bytearray) -> None:
+            frozen_class._encode(value, buffer)
+
+        # Will be overridden at finalization time.
+        frozen_class._encode = forward_encode
+
+        def forward_decode(stream: ByteStream) -> None:
+            return frozen_class._decode(stream)
+
+        frozen_class._decode = forward_decode
+
     def finalize(
         self,
         resolve_type_fn: Callable[[_spec.Type], TypeAdapter],
@@ -101,6 +115,18 @@ class StructAdapter(TypeAdapter):
             )
         )
 
+        num_slots_incl_removed: Final = max(
+            itertools.chain(
+                (f.field.number + 1 for f in fields),
+                (n + 1 for n in self.spec.removed_numbers),
+                [0],
+            )
+        )
+        self.num_slots_incl_removed = num_slots_incl_removed
+
+        slot_to_field: Final[list[_Field | None]] = [None] * self.num_slots_incl_removed
+        self.slot_to_field = slot_to_field
+
         # Aim to have dependencies finalized *before* the dependent. It's not always
         # possible, because there can be cyclic dependencies.
         # The function returned by the do_x_fn() method of a dependency is marginally
@@ -108,6 +134,7 @@ class StructAdapter(TypeAdapter):
         # this function is a "forwarding" function.
         for field in fields:
             field.type.finalize(resolve_type_fn)
+            self.slot_to_field[field.field.number] = field
 
         frozen_class = self.gen_class
         mutable_class = self.mutable_class
@@ -148,13 +175,29 @@ class StructAdapter(TypeAdapter):
             Any, _make_repr_fn(fields)
         )
 
-        frozen_class._tdj = _make_to_dense_json_fn(fields=fields)
+        frozen_class._tdj = _make_to_dense_json_fn(
+            fields=fields,
+            num_slots_incl_removed=num_slots_incl_removed,
+        )
         frozen_class._trj = _make_to_readable_json_fn(fields=fields)
+        frozen_class._encode = _make_encode_fn(
+            fields=fields,
+            num_slots_incl_removed=num_slots_incl_removed,
+        )
+        frozen_class._decode = _make_decode_fn(
+            frozen_class=frozen_class,
+            simple_class=simple_class,
+            fields=fields,
+            removed_numbers=self.spec.removed_numbers,
+            num_slots_incl_removed=num_slots_incl_removed,
+        )
+
         frozen_class._fj = _make_from_json_fn(
             frozen_class=frozen_class,
             simple_class=simple_class,
             fields=fields,
             removed_numbers=self.spec.removed_numbers,
+            num_slots_incl_removed=self.num_slots_incl_removed,
         )
 
         # Initialize DEFAULT.
@@ -187,7 +230,13 @@ class StructAdapter(TypeAdapter):
         )
 
     def is_not_default_expr(self, arg_expr: ExprLike, attr_expr: ExprLike) -> Expr:
-        return Expr.join(attr_expr, "._array_len")
+        return Expr.join(
+            "(",
+            attr_expr,
+            "._unrecognized or ",
+            attr_expr,
+            "._array_len)",
+        )
 
     def to_json_expr(
         self,
@@ -196,15 +245,29 @@ class StructAdapter(TypeAdapter):
     ) -> Expr:
         return Expr.join(in_expr, "._trj" if readable else "._tdj", "()")
 
-    def from_json_expr(self, json_expr: ExprLike) -> Expr:
+    def from_json_expr(
+        self, json_expr: ExprLike, keep_unrecognized_expr: ExprLike
+    ) -> Expr:
         fn_name = "_fj"
         # The _fj method may not have been added to the class yet.
         from_json_fn = getattr(self.gen_class, fn_name, None)
         if from_json_fn:
-            return Expr.join(Expr.local("_fj?", from_json_fn), "(", json_expr, ")")
+            return Expr.join(
+                Expr.local("_fj?", from_json_fn),
+                "(",
+                json_expr,
+                ", ",
+                keep_unrecognized_expr,
+                ")",
+            )
         else:
             return Expr.join(
-                Expr.local("_cls?", self.gen_class), f".{fn_name}(", json_expr, ")"
+                Expr.local("_cls?", self.gen_class),
+                f".{fn_name}(",
+                json_expr,
+                ", ",
+                keep_unrecognized_expr,
+                ")",
             )
 
     def get_type(self) -> reflection.Type:
@@ -236,6 +299,15 @@ class StructAdapter(TypeAdapter):
         for field in self.fields:
             field.type.register_records(registry)
 
+    def frozen_class_of_struct(self) -> type:
+        return self.gen_class
+
+    def encode_fn(self) -> Callable[[T, bytearray], None]:
+        return self.gen_class._encode
+
+    def decode_fn(self) -> Callable[[ByteStream], T]:
+        return self.gen_class._decode
+
 
 class _Frozen:
     def __setattr__(self, name: str, value: Any):
@@ -314,7 +386,7 @@ def _make_frozen_class_init_fn(
                 ")",
             )
         # Set the _unrecognized field.
-        builder.append_ln(obj_setattr, '(_self, "_unrecognized", ())')
+        builder.append_ln(obj_setattr, '(_self, "_unrecognized", None)')
         # Set array length.
         builder.append_ln(
             obj_setattr,
@@ -341,7 +413,7 @@ def _make_frozen_class_init_fn(
                 field.type.to_frozen_expr(attribute),
             )
         # Set the _unrecognized field.
-        builder.append_ln("_self._unrecognized = ()")
+        builder.append_ln("_self._unrecognized = None")
         # Set array length.
         builder.append_ln("_self._array_len = ", array_len_expr())
         # Change back the __class__.
@@ -375,7 +447,7 @@ def _make_mutable_class_init_fn(fields: Sequence[_Field]) -> Callable[..., None]
             " = ",
             attribute,
         )
-    builder.append_ln("_self._unrecognized = ()")
+    builder.append_ln("_self._unrecognized = None")
     return make_function(
         name="__init__",
         params=params,
@@ -524,26 +596,19 @@ def _make_to_frozen_fn(
     builder.append_ln("ret._unrecognized = self._unrecognized")
 
     def array_len_expr() -> Expr:
-        spans: list[LineSpanLike] = []
-        spans.append(
-            LineSpan.join(
-                f"{_get_num_slots(fields)} + ",
-                Expr.local("_len", len),
-                "(ret._unrecognized) if ret._unrecognized else ",
-            )
-        )
+        exprs: list[ExprLike] = []
         # Fields are sorted by number.
         for field in reversed(fields):
             attr_expr = f"ret.{field.field.attribute}"
-            spans.append(
+            exprs.append(
                 LineSpan.join(
                     f"{field.field.number + 1} if ",
                     field.type.is_not_default_expr(attr_expr, attr_expr),
                     " else ",
                 )
             )
-        spans.append("0")
-        return Expr.join(*spans)
+        exprs.append("0")
+        return Expr.join(*exprs)
 
     # Set the _unrecognized field.
     builder.append_ln("ret._unrecognized = self._unrecognized")
@@ -575,8 +640,14 @@ def _make_eq_fn(
     builder.append_ln("if other.__class__ is self.__class__:")
     operands: list[ExprLike]
     if fields:
-        attr: Callable[[_Field], str] = lambda f: f.field.attribute
-        operands = [Expr.join(f"self.{attr(f)} == other.{attr(f)}") for f in fields]
+
+        def get_attribute(f: _Field) -> str:
+            return f.field.attribute
+
+        operands = [
+            Expr.join(f"self.{get_attribute(f)} == other.{get_attribute(f)}")
+            for f in fields
+        ]
     else:
         operands = ["True"]
     builder.append_ln("  return ", Expr.join(*operands, separator=" and "))
@@ -639,9 +710,13 @@ def _make_repr_fn(fields: Sequence[_Field]) -> Callable[[Any], str]:
     )
 
 
-def _make_to_dense_json_fn(fields: Sequence[_Field]) -> Callable[[Any], Any]:
+def _make_to_dense_json_fn(
+    fields: Sequence[_Field], num_slots_incl_removed: int
+) -> Callable[[Any], Any]:
     builder = BodyBuilder()
-    builder.append_ln("l = self._array_len")
+    builder.append_ln(
+        "l = (self._unrecognized.adjusted_json_array_len if self._unrecognized else None) or self._array_len"
+    )
     builder.append_ln("ret = [0] * l")
     for field in fields:
         builder.append_ln(f"if l <= {field.field.number}:")
@@ -653,10 +728,8 @@ def _make_to_dense_json_fn(fields: Sequence[_Field]) -> Callable[[Any], Any]:
                 readable=False,
             ),
         )
-    num_slots = _get_num_slots(fields)
-    builder.append_ln(f"if l <= {num_slots}:")
-    builder.append_ln("  return ret")
-    builder.append_ln(f"ret[{num_slots}:] = self._unrecognized")
+    builder.append_ln(f"if {num_slots_incl_removed} < l:")
+    builder.append_ln(f"  ret[{num_slots_incl_removed}:] = self._unrecognized.json")
     builder.append_ln("return ret")
     return make_function(
         name="to_dense_json",
@@ -690,11 +763,60 @@ def _make_to_readable_json_fn(fields: Sequence[_Field]) -> Callable[[Any], Any]:
     )
 
 
+def _make_encode_fn(
+    fields: Sequence[_Field],
+    num_slots_incl_removed: int,
+) -> Callable[[Any, bytearray], None]:
+    builder = BodyBuilder()
+    builder.append_ln(
+        "l = (value._unrecognized.adjusted_bytes_array_len if ",
+        "value._unrecognized else None) or value._array_len",
+    )
+    builder.append_ln("if l < 4:")
+    builder.append_ln("  buffer.append(246 + l)")
+    builder.append_ln("else:")
+    builder.append_ln("  buffer.append(250)")
+    builder.append_ln(
+        "  ",
+        Expr.local("_encode_length_prefix", encode_length_prefix),
+        "(l, buffer)",
+    )
+    last_number = -1
+    for field in fields:
+        number = field.field.number
+        builder.append_ln(f"if l <= {number}:")
+        builder.append_ln("  return")
+        missing_slots = number - last_number - 1
+        if missing_slots >= 1:
+            # There are removed fields between last_number and number.
+            zeros = "\\0" * (missing_slots)
+            builder.append_ln(f"buffer.extend(b'{zeros}')")
+        builder.append_ln(
+            Expr.local("encode?", field.type.encode_fn()),
+            f"(value.{field.field.attribute}, buffer)",
+        )
+        last_number = number
+    builder.append_ln(f"if l <= {num_slots_incl_removed}:")
+    builder.append_ln("  return")
+    num_slots_excl_removed = last_number + 1
+    if num_slots_excl_removed < num_slots_incl_removed:
+        missing_slots = num_slots_incl_removed - num_slots_excl_removed
+        zeros = "\\0" * (missing_slots)
+        builder.append_ln(f"buffer.extend(b'{zeros}')")
+    builder.append_ln("buffer.extend(value._unrecognized.raw_bytes)")
+    return make_function(
+        name="encode",
+        params=["value, buffer"],
+        body=builder.build(),
+    )
+
+
 def _make_from_json_fn(
     frozen_class: type,
     simple_class: type,
     fields: Sequence[_Field],
     removed_numbers: tuple[int, ...],
+    num_slots_incl_removed: int,
 ) -> Callable[[Any], Any]:
     builder = BodyBuilder()
     builder.append_ln("if not json:")
@@ -714,29 +836,39 @@ def _make_from_json_fn(
         "(json)",
     )
     for field in fields:
-        name = field.field.name
         number = field.field.number
         item_expr = f"json[{number}]"
         builder.append_ln(
             f"  ret.{field.field.attribute} = ",
             field.type.default_expr(),
             f" if array_len <= {number} else ",
-            field.type.from_json_expr(item_expr),
+            field.type.from_json_expr(item_expr, "keep_unrecognized_fields"),
         )
-    num_slots = _get_num_slots(fields)
-    builder.append_ln(f"  if array_len <= {num_slots}:")
-    builder.append_ln("    ret._unrecognized = ()")
+    if fields:
+        num_slots_excl_removed = fields[-1].field.number + 1
+    else:
+        num_slots_excl_removed = 0
+    builder.append_ln(
+        f"  if array_len <= {num_slots_incl_removed} or not keep_unrecognized_fields:"
+    )
+    builder.append_ln("    ret._unrecognized = None")
     builder.append_ln(
         "    ret._array_len = ",
-        _adjust_array_len_expr("array_len", removed_numbers),
+        _adjust_array_len_expr(
+            "array_len",
+            removed_numbers=removed_numbers,
+            num_slots_excl_removed=num_slots_excl_removed,
+        ),
     )
     builder.append_ln("  else:")
     builder.append_ln(
         "    ret._unrecognized = ",
+        Expr.local("UnrecognizedFields", _UnrecognizedFields.from_json),
+        "(json=",
         Expr.local("deepcopy", copy.deepcopy),
-        f"(json[{num_slots}:])",
+        f"(json[{num_slots_incl_removed}:]), adjusted_json_array_len=array_len)",
     )
-    builder.append_ln("    ret._array_len = array_len")
+    builder.append_ln(f"    ret._array_len = {num_slots_excl_removed}")
 
     builder.append_ln("else:")
     builder.append_ln("  array_len = 0")
@@ -748,12 +880,12 @@ def _make_from_json_fn(
         builder.append_ln(f"    array_len = {field.field.number + 1}")
         builder.append_ln(
             f"    {lvalue} = ",
-            field.type.from_json_expr(f'json["{name}"]'),
+            field.type.from_json_expr(f'json["{name}"]', "keep_unrecognized_fields"),
         )
         builder.append_ln("  else:")
         builder.append_ln(f"    {lvalue} = ", field.type.default_expr())
     # Drop unrecognized fields in readable mode.
-    builder.append_ln("  ret._unrecognized = ()")
+    builder.append_ln("  ret._unrecognized = None")
     builder.append_ln("  ret._array_len = array_len")
 
     builder.append_ln("ret.__class__ = ", Expr.local("Frozen", frozen_class))
@@ -761,60 +893,121 @@ def _make_from_json_fn(
 
     return make_function(
         name="from_json",
-        params=["json"],
+        params=["json", "keep_unrecognized_fields"],
+        body=builder.build(),
+    )
+
+
+def _make_decode_fn(
+    frozen_class: type,
+    simple_class: type,
+    fields: Sequence[_Field],
+    removed_numbers: tuple[int, ...],
+    num_slots_incl_removed: int,
+) -> Callable[[ByteStream], Any]:
+    builder = BodyBuilder()
+    builder.append_ln("wire = stream.buffer[stream.position]")
+    builder.append_ln("stream.position += 1")
+    builder.append_ln("if wire in (0, 246):")
+    builder.append_ln("  return ", Expr.local("DEFAULT", frozen_class.DEFAULT))
+    builder.append_ln("elif wire == 250:")
+    builder.append_ln(
+        "  array_len = ",
+        Expr.local("decode_int64", decode_int64),
+        "(stream)",
+    )
+    builder.append_ln("else:")
+    builder.append_ln("  array_len = wire - 246")
+    # Create an instance of the simple class. We'll later change its __class__ attr.
+    builder.append_ln(
+        "ret = ",
+        Expr.local("Simple", simple_class),
+        "()",
+    )
+    last_number = -1
+    for field in fields:
+        number = field.field.number
+        for removed_number in range(last_number + 1, number):
+            builder.append_ln(f"if {removed_number} < array_len:")
+            builder.append_ln(
+                "  ",
+                Expr.local("decode_unused", decode_unused),
+                "(stream)",
+            )
+        builder.append_ln(
+            f"ret.{field.field.attribute} = ",
+            field.type.default_expr(),
+            f" if array_len <= {number} else ",
+            Expr.local("decode?", field.type.decode_fn()),
+            "(stream)",
+        )
+        last_number = number
+    num_slots_excl_removed = last_number + 1
+
+    builder.append_ln(f"if array_len > {num_slots_excl_removed}:")
+    builder.append_ln(
+        f"  if array_len > {num_slots_incl_removed} and keep_unrecognized_fields:"
+    )
+    for _ in range(num_slots_incl_removed - num_slots_excl_removed):
+        builder.append_ln(Expr.local("    decode_unused", decode_unused), "(stream)")
+    builder.append_ln("    start_offset = stream.position")
+    builder.append_ln("    for _ in range(array_len - {num_slots_incl_removed}):")
+    builder.append_ln("      ", Expr.local("decode_unused", decode_unused), "(stream)")
+    builder.append_ln("    end_offset = stream.position")
+    builder.append_ln(
+        "    ret._unrecognized = ",
+        Expr.local("UnrecognizedFields", _UnrecognizedFields.from_bytes),
+        "(raw_bytes=stream.buffer[start_offset:end_offset], adjusted_bytes_array_len=array_len)",
+    )
+    builder.append_ln("  else:")
+    builder.append_ln(f"    for _ in range(array_len - {num_slots_excl_removed}):")
+    builder.append_ln("      ", Expr.local("decode_unused", decode_unused), "(stream)")
+    builder.append_ln("    ret._unrecognized = None")
+    builder.append_ln("else:")
+    builder.append_ln("  ret._unrecognized = None")
+
+    builder.append_ln(
+        "  ret._array_len = ",
+        _adjust_array_len_expr(
+            "array_len",
+            removed_numbers=removed_numbers,
+            num_slots_excl_removed=num_slots_excl_removed,
+        ),
+    )
+
+    builder.append_ln(
+        "ret.__class__ = ",
+        Expr.local("Frozen", frozen_class),
+    )
+    builder.append_ln("return ret")
+    return make_function(
+        name="decode",
+        params=["stream"],
         body=builder.build(),
     )
 
 
-def _adjust_array_len_expr(var: str, removed_numbers: tuple[int, ...]) -> str:
+def _adjust_array_len_expr(
+    var: str,
+    removed_numbers: tuple[int, ...],
+    num_slots_excl_removed: int,
+) -> str:
     """
     When parsing a dense JSON or decoding a binary string, we can reuse the array length
     in the decoded struct, but we need to account for possibly newly-removed fields. The
     last field of the adjusted array length cannot be a removed field.
-
-    Let's imagine that field number 3 was removed from a struct.
-    This function would return the following expression:
-        array_len if array_len <= 3 else 3 if array_len == 4 else array_len
     """
-
-    @dataclass
-    class _RemovedSpan:
-        """Sequence of consecutive removed fields."""
-
-        # Number of the first removed field.
-        begin: int = 0
-        # Number after the last removed field.
-        end: int = 0
-
-    def get_removed_spans() -> list[_RemovedSpan]:
-        ret: list[_RemovedSpan] = []
-        for number in sorted(removed_numbers):
-            last = ret[-1] if ret else None
-            if last and last.end == number:
-                last.end = number + 1
-            else:
-                ret.append(_RemovedSpan(number, number + 1))
-        return ret
-
-    removed_spans = get_removed_spans()
-
-    ret = ""
-    lower_bound = 0
-    for s in removed_spans:
-        if s.begin == lower_bound:
-            ret += f"{s.begin} if {var} <= {s.end} else "
-        elif s.end == s.begin + 1:
-            # Similar to the expression in 'else' but uses '==' instead of '<='
-            ret += (
-                f"{var} if {var} <= {s.begin} else {s.begin} if {var} == {s.end} else "
-            )
+    removed_numbers_set = set(removed_numbers)
+    table: list[int] = [0]
+    last_len = 0
+    for i in range(0, num_slots_excl_removed - 1):
+        if i in removed_numbers_set:
+            table.append(last_len)
         else:
-            ret += (
-                f"{var} if {var} <= {s.begin} else {s.begin} if {var} <= {s.end} else "
-            )
-        lower_bound = s.end + 1
-    ret += var
-    return ret
+            table.append(i + 1)
+            last_len = i + 1
+    table_tuple = tuple(table)
+    return f"{table_tuple}[{var}] if {var} < {num_slots_excl_removed} else {num_slots_excl_removed}"
 
 
 def _init_default(default: Any, fields: Sequence[_Field]) -> None:
@@ -826,7 +1019,7 @@ def _init_default(default: Any, fields: Sequence[_Field]) -> None:
             body=(Line.join("return ", field.type.default_expr()),),
         )
         object.__setattr__(default, attribute, get_field_default())
-    object.__setattr__(default, "_unrecognized", ())
+    object.__setattr__(default, "_unrecognized", None)
     object.__setattr__(default, "_array_len", 0)
 
 
@@ -874,5 +1067,38 @@ def _name_private_is_frozen_attr(record_id: str) -> str:
     return f"_is_{record_name}_{hex_hash}"
 
 
-def _get_num_slots(fields: Sequence[_Field]) -> int:
-    return (fields[-1].field.number + 1) if fields else 0
+@dataclass(frozen=True)
+class _UnrecognizedFields:
+    __slots__ = (
+        "json",
+        "raw_bytes",
+        "adjusted_json_array_len",
+        "adjusted_bytes_array_len",
+    )
+
+    json: list[Any] | None
+    raw_bytes: bytes | None
+    adjusted_json_array_len: int | None
+    adjusted_bytes_array_len: int | None
+
+    @staticmethod
+    def from_json(
+        json: list[Any], adjusted_json_array_len: int
+    ) -> "_UnrecognizedFields":
+        return _UnrecognizedFields(
+            json=json,
+            raw_bytes=None,
+            adjusted_json_array_len=adjusted_json_array_len,
+            adjusted_bytes_array_len=None,
+        )
+
+    @staticmethod
+    def from_bytes(
+        raw_bytes: bytes, adjusted_bytes_array_len: int
+    ) -> "_UnrecognizedFields":
+        return _UnrecognizedFields(
+            json=None,
+            raw_bytes=raw_bytes,
+            adjusted_json_array_len=None,
+            adjusted_bytes_array_len=adjusted_bytes_array_len,
+        )

soia/_impl/timestamp.py

@@ -146,11 +146,10 @@ class Timestamp:
             dt = self.to_datetime_or_raise()
         except Exception:
             return ""
-        if dt:
-            ret = dt.isoformat()
-            bad_suffix = "+00:00"
-            if ret.endswith(bad_suffix):
-                ret = ret[0 : -len(bad_suffix)] + "Z"
+        ret = dt.isoformat()
+        bad_suffix = "+00:00"
+        if ret.endswith(bad_suffix):
+            ret = ret[0 : -len(bad_suffix)] + "Z"
         return ret