Merge pull request #525 from pikers/msgspec_struct_updates

`msgspec` struct derivative updates

Author: ebisu4

piker/data/types.py

@@ -1,5 +1,7 @@
 # piker: trading gear for hackers
-# Copyright (C) Guillermo Rodriguez (in stewardship for piker0)
+# Copyright (C) (in stewardship for pikers)
+#  - Tyler Goodlet
+#  - Guillermo Rodriguez
 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as published by
@@ -14,18 +16,22 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-"""
-Built-in (extension) types.
+'''
+Extensions to built-in or (heavily used but 3rd party) friend-lib
+types.
 
-"""
-import sys
+'''
 from pprint import pformat
 
-import msgspec
+from msgspec import (
+    msgpack,
+    Struct,
+    structs,
+)
 
 
 class Struct(
-    msgspec.Struct,
+    Struct,
 
     # https://jcristharif.com/msgspec/structs.html#tagged-unions
     # tag='pikerstruct',
@@ -36,22 +42,14 @@ class Struct(
 
     '''
     def to_dict(self) -> dict:
-        return {
-            f: getattr(self, f)
-            for f in self.__struct_fields__
-        }
-
-    # Lul, doesn't seem to work that well..
-    # def __repr__(self):
-    #     # only turn on pprint when we detect a python REPL
-    #     # at runtime B)
-    #     if (
-    #         hasattr(sys, 'ps1')
-    #         # TODO: check if we're in pdb
-    #     ):
-    #         return self.pformat()
-
-    #     return super().__repr__()
+        '''
+        Like it sounds.. direct delegation to:
+        https://jcristharif.com/msgspec/api.html#msgspec.structs.asdict
+
+        TODO: probably just drop this method since it's now a built-int method?
+
+        '''
+        return structs.asdict(self)
 
     def pformat(self) -> str:
         return f'Struct({pformat(self.to_dict())})'
@@ -60,30 +58,47 @@ def copy(
         self,
         update: dict | None = None,
 
-    ) -> msgspec.Struct:
+    ) -> Struct:
         '''
-        Validate-typecast all self defined fields, return a copy of us
-        with all such fields.
+        Validate-typecast all self defined fields, return a copy of
+        us with all such fields.
 
-        This is kinda like the default behaviour in `pydantic.BaseModel`.
+        NOTE: This is kinda like the default behaviour in
+        `pydantic.BaseModel` except a copy of the object is
+        returned making it compat with `frozen=True`.
 
         '''
         if update:
             for k, v in update.items():
                 setattr(self, k, v)
 
-        # roundtrip serialize to validate
-        return msgspec.msgpack.Decoder(
-            type=type(self)
-        ).decode(
-            msgspec.msgpack.Encoder().encode(self)
+        # NOTE: roundtrip serialize to validate
+        # - enode to msgpack binary format,
+        # - decode that back to a struct.
+        return msgpack.Decoder(type=type(self)).decode(
+            msgpack.Encoder().encode(self)
         )
 
-    # NOTE XXX: this won't work on frozen types!
-    # use ``.copy()`` above in such cases.
     def typecast(
         self,
-        # fields: list[str] | None = None,
+
+        # TODO: allow only casting a named subset?
+        # fields: set[str] | None = None,
+
     ) -> None:
-        for fname, ftype in self.__annotations__.items():
-            setattr(self, fname, ftype(getattr(self, fname)))
+        '''
+        Cast all fields using their declared type annotations
+        (kinda like what `pydantic` does by default).
+
+        NOTE: this of course won't work on frozen types, use
+        ``.copy()`` above in such cases.
+
+        '''
+        # https://jcristharif.com/msgspec/api.html#msgspec.structs.fields
+        fi: structs.FieldInfo
+        for fi in structs.fields(self):
+            setattr(
+                self,
+                fi.name,
+                fi.type(getattr(self, fi.name)),
+            )