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

soia/_impl/enums.py

@@ -1,15 +1,16 @@
 import copy
 from collections.abc import Callable, Sequence
 from dataclasses import FrozenInstanceError, dataclass
-from typing import Any, Final, Union
+from typing import Any, Final, Generic, Union
 
 from soia import _spec, reflection
+from soia._impl.binary import decode_int64, decode_unused, encode_int64
 from soia._impl.function_maker import BodyBuilder, Expr, ExprLike, Line, make_function
 from soia._impl.repr import repr_impl
-from soia._impl.type_adapter import TypeAdapter
+from soia._impl.type_adapter import T, ByteStream, TypeAdapter
 
 
-class EnumAdapter(TypeAdapter):
+class EnumAdapter(Generic[T], TypeAdapter[T]):
     __slots__ = (
         "spec",
         "gen_class",
@@ -30,6 +31,12 @@ class EnumAdapter(TypeAdapter):
         self.spec = spec
         base_class = self.gen_class = _make_base_class(spec)
 
+        def forward_decode(stream: ByteStream) -> T:
+            return base_class._decode(stream)
+
+        # Will be overridden at finalization time.
+        base_class._decode = forward_decode
+
         private_is_enum_attr = _name_private_is_enum_attr(spec.id)
         self.private_is_enum_attr = private_is_enum_attr
         setattr(base_class, private_is_enum_attr, True)
@@ -73,12 +80,28 @@ class EnumAdapter(TypeAdapter):
         for value_field in value_fields:
             wrap_fn = _make_wrap_fn(value_field)
             setattr(base_class, f"wrap_{value_field.spec.name}", wrap_fn)
+            # Check if the field type is a struct type.
+            field_type = resolve_type_fn(value_field.spec.type)
+            frozen_class = field_type.frozen_class_of_struct()
+            if frozen_class:
+                create_fn = _make_create_fn(wrap_fn, frozen_class)
+                setattr(base_class, f"create_{value_field.spec.name}", create_fn)
+
+        unrecognized_class = _make_unrecognized_class(base_class)
 
         base_class._fj = _make_from_json_fn(
             self.all_constant_fields,
             value_fields,
             set(self.spec.removed_numbers),
-            base_class,
+            base_class=base_class,
+            unrecognized_class=unrecognized_class,
+        )
+        base_class._decode = _make_decode_fn(
+            self.all_constant_fields,
+            value_fields,
+            set(self.spec.removed_numbers),
+            base_class=base_class,
+            unrecognized_class=unrecognized_class,
         )
 
         # Mark finalization as done.
@@ -111,16 +134,36 @@ class EnumAdapter(TypeAdapter):
     def to_json_expr(self, in_expr: ExprLike, readable: bool) -> Expr:
         return Expr.join(in_expr, "._rj" if readable else "._dj")
 
-    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"
         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 encode_fn(self) -> Callable[[T, bytearray], None]:
+        return _encode_impl
+
+    def decode_fn(self) -> Callable[[ByteStream], T]:
+        return self.gen_class._decode
+
     def get_type(self) -> reflection.Type:
         return reflection.RecordType(
             kind="record",
@@ -143,8 +186,10 @@ class EnumAdapter(TypeAdapter):
                     number=field.number,
                     type=None,
                 )
-                for field in self.all_constant_fields if field.number != 0
-            ) + tuple(
+                for field in self.all_constant_fields
+                if field.number != 0
+            )
+            + tuple(
                 reflection.Field(
                     name=field.spec.name,
                     number=field.spec.number,
@@ -157,6 +202,9 @@ class EnumAdapter(TypeAdapter):
         for field in self.value_fields:
             field.field_type.register_records(registry)
 
+    def frozen_class_of_struct(self) -> type | None:
+        return None
+
 
 def _make_base_class(spec: _spec.Enum) -> type:
     record_hash = hash(spec.id)
@@ -174,6 +222,9 @@ def _make_base_class(spec: _spec.Enum) -> type:
         def union(self) -> Any:
             return self
 
+        def __bool__(self) -> bool:
+            return self.kind != "?"
+
         def __setattr__(self, name: str, value: Any):
             raise FrozenInstanceError(self.__class__.__qualname__)
 
@@ -195,6 +246,9 @@ def _make_base_class(spec: _spec.Enum) -> type:
 
 
 def _make_constant_class(base_class: type, spec: _spec.ConstantField) -> type:
+    byte_array = bytearray()
+    encode_int64(spec.number, byte_array)
+
     class Constant(base_class):
         __slots__ = ()
 
@@ -206,6 +260,7 @@ def _make_constant_class(base_class: type, spec: _spec.ConstantField) -> type:
         _rj: Final[str] = spec.name
         # has value
         _hv: Final[bool] = False
+        _bytes: Final[bytes | None] = bytes(byte_array)
 
         def __init__(self):
             # Do not call super().__init__().
@@ -225,20 +280,22 @@ def _make_unrecognized_class(base_class: type) -> type:
     """
 
     class Unrecognized(base_class):
-        __slots__ = ("_dj",)
+        __slots__ = ("_dj", "_bytes")
 
         kind: Final[str] = "?"
         _number: Final[int] = 0
         # dense JSON
-        _dj: Any
+        _dj: list[Any] | int
+        _bytes: bytes
         # readable JSON
         _rj: Final[str] = "?"
         # has value
         _hv: Final[bool] = False
 
-        def __init__(self, dj: Any):
+        def __init__(self, dj: list[Any] | int, bytes: bytes):
             # Do not call super().__init__().
             object.__setattr__(self, "_dj", copy.deepcopy(dj))
+            object.__setattr__(self, "_bytes", bytes)
             object.__setattr__(self, "value", None)
 
         def __repr__(self) -> str:
@@ -261,6 +318,7 @@ def _make_value_class(
         _number: Final[int] = number
         # has value
         _hv: Final[bool] = True
+        _bytes: Final[None] = None
 
         def __init__(self):
             # Do not call super().__init__().
@@ -305,9 +363,38 @@ def _make_value_class(
         )
     )
 
+    bytes_prefix = bytearray()
+    if number in range(1, 5):
+        bytes_prefix.append(250 + number)
+    else:
+        bytes_prefix.append(248)
+        encode_int64(number, bytes_prefix)
+
+    ret._enc = make_function(
+        name="encode",
+        params=["self", "buffer"],
+        body=[
+            f"buffer.extend({bytes_prefix})",
+            Line.join(
+                Expr.local("encode_value", field_type.encode_fn()),
+                "(self.value, buffer)",
+            ),
+        ],
+    )
+
     return ret
 
 
+def _encode_impl(
+    value: Any,
+    buffer: bytearray,
+) -> None:
+    if value._bytes:
+        buffer.extend(value._bytes)
+    else:
+        value._enc(buffer)
+
+
 @dataclass(frozen=True)
 class _ValueField:
     spec: _spec.ValueField
@@ -344,16 +431,23 @@ def _make_wrap_fn(field: _ValueField) -> Callable[[Any], Any]:
     )
 
 
+def _make_create_fn(wrap_fn: Callable[[Any], Any], frozen_class: type) -> Callable:
+    def create(**kwargs):
+        return wrap_fn(frozen_class(**kwargs))
+
+    return create
+
+
 def _make_from_json_fn(
     constant_fields: Sequence[_spec.ConstantField],
     value_fields: Sequence[_ValueField],
     removed_numbers: set[int],
     base_class: type,
+    unrecognized_class: type,
 ) -> Callable[[Any], Any]:
-    unrecognized_class = _make_unrecognized_class(base_class)
     unrecognized_class_local = Expr.local("Unrecognized", unrecognized_class)
     obj_setattr_local = Expr.local("obj_settatr", object.__setattr__)
-    removed_numbers_local = Expr.local("removed_numbers", removed_numbers)
+    removed_numbers_tuple = tuple(sorted(removed_numbers))
 
     key_to_constant: dict[Union[int, str], Any] = {}
     for field in constant_fields:
@@ -364,15 +458,13 @@ def _make_from_json_fn(
     unknown_constant = key_to_constant[0]
     unknown_constant_local = Expr.local("unknown_constant", unknown_constant)
 
-    numbers: list[int] = []
-    names: list[str] = []
-    key_to_field: dict[Union[int, str], _ValueField] = {}
+    number_to_value_field: dict[int, _ValueField] = {}
+    name_to_value_field: dict[str, _ValueField] = {}
     for field in value_fields:
-        numbers.append(field.spec.number)
-        names.append(field.spec.name)
-        key_to_field[field.spec.number] = field
-        key_to_field[field.spec.name] = field
-    value_keys_local = Expr.local("value_keys", set(key_to_field.keys()))
+        number_to_value_field[field.spec.number] = field
+        name_to_value_field[field.spec.name] = field
+    value_field_numbers = tuple(sorted(number_to_value_field.keys()))
+    value_field_names = tuple(sorted(name_to_value_field.keys()))
 
     builder = BodyBuilder()
     # The reason why we wrap the function inside a 'while' is explained below.
@@ -389,16 +481,20 @@ def _make_from_json_fn(
         builder.append_ln("      return ", key_to_constant_local, "[json]")
         builder.append_ln("    except:")
         if removed_numbers:
-            builder.append_ln("      if json in ", removed_numbers_local, ":")
+            builder.append_ln(
+                f"      if json in {removed_numbers_tuple} or not keep_unrecognized_fields:"
+            )
             builder.append_ln("        return ", unknown_constant_local)
-        builder.append_ln("      return ", unrecognized_class_local, "(json)")
+        builder.append_ln("      return ", unrecognized_class_local, "(json, b'\\0')")
 
-    def append_number_branches(numbers: list[int], indent: str) -> None:
+    def append_number_branches(numbers: Sequence[int], indent: str) -> None:
         if len(numbers) == 1:
             number = numbers[0]
-            field = key_to_field[number]
+            field = number_to_value_field[number]
             value_class_local = Expr.local("cls?", field.value_class)
-            value_expr = field.field_type.from_json_expr("json[1]")
+            value_expr = field.field_type.from_json_expr(
+                "json[1]", "keep_unrecognized_fields"
+            )
             builder.append_ln(f"{indent}ret = ", value_class_local, "()")
             builder.append_ln(
                 indent, obj_setattr_local, '(ret, "value", ', value_expr, ")"
@@ -420,19 +516,20 @@ def _make_from_json_fn(
     if not value_fields:
         # The field was either removed or is an unrecognized field.
         if removed_numbers:
-            builder.append_ln("    if number in ", removed_numbers_local, ":")
+            builder.append_ln(
+                f"    if number in {removed_numbers_tuple} or not keep_unrecognized_fields:"
+            )
             builder.append_ln("      return ", unknown_constant_local)
-        builder.append_ln("    return ", unrecognized_class_local, "(json)")
+        builder.append_ln("    return ", unrecognized_class_local, "(json, b'\\0')")
     else:
-        if len(value_fields) == 1:
-            builder.append_ln(f"    if number != {value_fields[0].spec.number}:")
-        else:
-            builder.append_ln("    if number not in ", value_keys_local, ":")
+        builder.append_ln(f"    if number not in {value_field_numbers}:")
         if removed_numbers:
-            builder.append_ln("      if number in ", removed_numbers_local, ":")
+            builder.append_ln(
+                f"      if number in {removed_numbers_tuple} or not keep_unrecognized_fields:"
+            )
             builder.append_ln("        return ", unknown_constant_local)
-        builder.append_ln("      return ", unrecognized_class_local, "(json)")
-        append_number_branches(sorted(numbers), "    ")
+        builder.append_ln("      return ", unrecognized_class_local, "(json, b'\\0')")
+        append_number_branches(value_field_numbers, "    ")
 
     # READABLE FORMAT
     if len(constant_fields) == 1:
@@ -446,12 +543,14 @@ def _make_from_json_fn(
         # In readable mode, drop unrecognized values and use UNKNOWN instead.
         builder.append_ln("      return ", unknown_constant_local)
 
-    def append_name_branches(names: list[str], indent: str) -> None:
+    def append_name_branches(names: Sequence[str], indent: str) -> None:
         if len(names) == 1:
             name = names[0]
-            field = key_to_field[name]
+            field = name_to_value_field[name]
             value_class_local = Expr.local("cls?", field.value_class)
-            value_expr = field.field_type.from_json_expr("json['value']")
+            value_expr = field.field_type.from_json_expr(
+                "json['value']", "keep_unrecognized_fields"
+            )
             builder.append_ln(f"{indent}ret = ", value_class_local, "()")
             builder.append_ln(
                 indent, obj_setattr_local, '(ret, "value", ', value_expr, ")"
@@ -472,10 +571,10 @@ def _make_from_json_fn(
         builder.append_ln("    return ", unknown_constant_local)
     else:
         builder.append_ln("    kind = json['kind']")
-        builder.append_ln("    if kind not in ", value_keys_local, ":")
+        builder.append_ln(f"    if kind not in {value_field_names}:")
         builder.append_ln("      return ", unknown_constant_local)
         builder.append_ln("    else:")
-        append_name_branches(sorted(names), "      ")
+        append_name_branches(value_field_names, "      ")
 
     # In the unlikely event that json.loads() returns an instance of a subclass of int.
     builder.append_ln("  elif isinstance(json, int):")
@@ -487,7 +586,112 @@ 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(
+    constant_fields: Sequence[_spec.ConstantField],
+    value_fields: Sequence[_ValueField],
+    removed_numbers: set[int],
+    base_class: type,
+    unrecognized_class: type,
+) -> Callable[[ByteStream], Any]:
+    unrecognized_class_local = Expr.local("Unrecognized", unrecognized_class)
+    obj_setattr_local = Expr.local("obj_settatr", object.__setattr__)
+
+    number_to_constant: dict[int, Any] = {}
+    for field in constant_fields:
+        constant = getattr(base_class, field.attribute)
+        number_to_constant[field.number] = constant
+    number_to_constant_local = Expr.local("number_to_constant", number_to_constant)
+    removed_numbers_tuple = tuple(sorted(removed_numbers))
+    unknown_constant = number_to_constant[0]
+    unknown_constant_local = Expr.local("unknown_constant", unknown_constant)
+
+    number_to_value_field: dict[int, _ValueField] = {}
+    for field in value_fields:
+        number_to_value_field[field.spec.number] = field
+    value_field_numbers = tuple(sorted(number_to_value_field.keys()))
+
+    builder = BodyBuilder()
+    builder.append_ln("start_offset = stream.position")
+    builder.append_ln("wire = stream.buffer[start_offset]")
+    builder.append_ln("if wire <= 238:")
+    # A number
+    builder.append_ln("  if wire < 232:")
+    builder.append_ln("    stream.position += 1")
+    builder.append_ln("    number = wire")
+    builder.append_ln("  else:")
+    builder.append_ln(
+        "    number = ", Expr.local("decode_int64", decode_int64), "(stream)"
+    )
+    builder.append_ln("  try:")
+    builder.append_ln("    return ", number_to_constant_local, "[number]")
+    builder.append_ln("  except:")
+    builder.append_ln("    ", Expr.local("decode_unused", decode_unused), "(stream)")
+    if removed_numbers:
+        builder.append_ln(
+            f"    if number in {removed_numbers_tuple} or not keep_unrecognized_fields:"
+        )
+        builder.append_ln("      return ", unknown_constant_local)
+    builder.append_ln("    bytes = stream.buffer[start_offset:stream.position]")
+    builder.append_ln("    return ", unrecognized_class_local, "(0, bytes)")
+    # An array of 2
+    builder.append_ln("stream.position += 1")
+    builder.append_ln("if wire == 248:")
+    builder.append_ln(
+        "  number = ", Expr.local("decode_int64", decode_int64), "(stream)"
+    )
+    builder.append_ln("else:")
+    builder.append_ln("  number = wire - 250")
+
+    def append_number_branches(numbers: Sequence[int], indent: str) -> None:
+        if len(numbers) == 1:
+            number = numbers[0]
+            field = number_to_value_field[number]
+            value_class_local = Expr.local("cls?", field.value_class)
+            decode_local = Expr.local("decode?", field.field_type.decode_fn())
+            value_expr = Expr.join(decode_local, "(stream)")
+            builder.append_ln(f"{indent}ret = ", value_class_local, "()")
+            builder.append_ln(
+                indent, obj_setattr_local, '(ret, "value", ', value_expr, ")"
+            )
+            builder.append_ln(f"{indent}return ret")
+        else:
+            indented = f"  {indent}"
+            mid_index = int(len(numbers) / 2)
+            mid_number = numbers[mid_index - 1]
+            operator = "==" if mid_index == 1 else "<="
+            builder.append_ln(f"{indent}if number {operator} {mid_number}:")
+            append_number_branches(numbers[0:mid_index], indented)
+            builder.append_ln(f"{indent}else:")
+            append_number_branches(numbers[mid_index:], indented)
+
+    if not value_fields:
+        # The field was either removed or is an unrecognized field.
+        if removed_numbers:
+            builder.append_ln(
+                f"if number in {removed_numbers_tuple} or not keep_unrecognized_fields:"
+            )
+            builder.append_ln("  return ", unknown_constant_local)
+        builder.append_ln("bytes = stream.buffer[start_offset:stream.position]")
+        builder.append_ln("return ", unrecognized_class_local, "(0, bytes)")
+    else:
+        builder.append_ln(f"if number not in {value_field_numbers}:")
+        if removed_numbers:
+            builder.append_ln(
+                f"  if number in {removed_numbers_tuple} or not keep_unrecognized_fields:"
+            )
+            builder.append_ln("    return ", unknown_constant_local)
+        builder.append_ln("  bytes = stream.buffer[start_offset:stream.position]")
+        builder.append_ln("  return ", unrecognized_class_local, "(0, bytes)")
+        append_number_branches(value_field_numbers, "")
+
+    return make_function(
+        name="decode",
+        params=["stream"],
         body=builder.build(),
     )
 

soia/_impl/optionals.py

@@ -1,24 +1,22 @@
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TypeVar
+from functools import cached_property
+from typing import Generic
 from weakref import WeakValueDictionary
 
-from soia._impl.function_maker import Expr, ExprLike
-from soia._impl.type_adapter import TypeAdapter
-
 from soia import _spec, reflection
-
-Other = TypeVar("Other")
+from soia._impl.function_maker import Expr, ExprLike
+from soia._impl.type_adapter import T, ByteStream, TypeAdapter
 
 
-def get_optional_adapter(other_adapter: TypeAdapter) -> TypeAdapter:
+def get_optional_adapter(other_adapter: TypeAdapter[T]) -> TypeAdapter[T | None]:
     return _other_adapter_to_optional_adapter.setdefault(
         other_adapter, _OptionalAdapter(other_adapter)
     )
 
 
 @dataclass(frozen=True)
-class _OptionalAdapter(TypeAdapter):
+class _OptionalAdapter(Generic[T], TypeAdapter[T | None]):
     __slots__ = ("other_adapter",)
 
     other_adapter: TypeAdapter
@@ -50,8 +48,12 @@ class _OptionalAdapter(TypeAdapter):
             ")",
         )
 
-    def from_json_expr(self, json_expr: ExprLike) -> ExprLike:
-        other_from_json = self.other_adapter.from_json_expr(json_expr)
+    def from_json_expr(
+        self, json_expr: ExprLike, keep_unrecognized_expr: ExprLike
+    ) -> ExprLike:
+        other_from_json = self.other_adapter.from_json_expr(
+            json_expr, keep_unrecognized_expr
+        )
         if other_from_json == json_expr:
             return json_expr
         return Expr.join(
@@ -62,6 +64,42 @@ class _OptionalAdapter(TypeAdapter):
             ")",
         )
 
+    @cached_property
+    def encode_fn_impl(self) -> Callable[[T | None, bytearray], None]:
+        encode_value = self.other_adapter.encode_fn()
+
+        def encode(
+            value: T | None,
+            buffer: bytearray,
+        ) -> None:
+            if value is None:
+                buffer.append(255)
+            else:
+                encode_value(value, buffer)
+
+        return encode
+
+    def encode_fn(self) -> Callable[[T | None, bytearray], None]:
+        return self.encode_fn_impl
+
+    @cached_property
+    def decode_fn_impl(self) -> Callable[[ByteStream], T | None]:
+        decode_value = self.other_adapter.decode_fn()
+
+        def decode(
+            stream: ByteStream,
+        ) -> T | None:
+            if stream.buffer[stream.position] == 255:
+                stream.position += 1
+                return None
+            else:
+                return decode_value(stream)
+
+        return decode
+
+    def decode_fn(self) -> Callable[[ByteStream], T | None]:
+        return self.decode_fn_impl
+
     def finalize(
         self,
         resolve_type_fn: Callable[[_spec.Type], "TypeAdapter"],
@@ -80,6 +118,9 @@ class _OptionalAdapter(TypeAdapter):
     ) -> None:
         self.other_adapter.register_records(registry)
 
+    def frozen_class_of_struct(self) -> type | None:
+        return None
+
 
 _other_adapter_to_optional_adapter: WeakValueDictionary[TypeAdapter, TypeAdapter] = (
     WeakValueDictionary()