More ergonomic pattern for referencing StructDef or HashVariantDef in the protocol tree

[johnlevidy]

Today, there's not really a nice way to do this, for the example protocol:

class Bakery(Protocol):    
    PacketVariant = HashVariant[lu32]([V1.Message, V2.Message, V3.Message, V4.Message])    
    
    Packet = Struct(payload=PacketVariant

Bakery.Packet.fields['payload']

gets me the HashVariantDef

I'd like to be able to do something like

Bakery.Packet.payload.

The difficulty will probably lie in what we end up doing for variants, since today a HashVariantDef just stores a list of the underlying types.

wat(BakeryActor.Produces[0].fields['payload'].hash_types[0])

It's probable for all Variants I can have some simple getter function that gets by name? The problem is names can collide. Above, only the QName stores the full name. The type itself just calls itself "Message"

[johnlevidy]

Steps to reproduce in REPL

from test_types.bakery.v1 import V1, V2
from test_types.bakery import Bakery
import wat
wat(Bakery.Packet.fields['payload'])
wat(Bakery.Packet.fields['payload'].get_variant_by_type(type(V1.Message)))
wat(Bakery.Packet.fields['payload'].get_variant_by_type(type(V2.Message)))

With some minor changes this is working ^^

diff --git a/python/tako/core/types.py b/python/tako/core/types.py
index 20743d5..3d263c1 100644
--- a/python/tako/core/types.py
+++ b/python/tako/core/types.py
@@ -269,6 +269,13 @@ class HashVariantDef(RootType):
     tag_type: Int
     hash_types: t.List[StructDef]
 
+    def get_variant_by_type(self, match_type: t.Type) -> Optional[T]:
+        for h in self.hash_types:
+            print(h.name)
+            if isinstance(h, match_type):
+                return h
+        return None
+
     def accept(self, visitor: TypeVisitor[T]) -> T:
         return visitor.visit_hash_variant_def(self)
 

This at least does what we want it just needs to be a lot more ergonomic to reference subtrees in the protocol

[johnlevidy]

index 20743d5..899ace6 100644
--- a/python/tako/core/types.py
+++ b/python/tako/core/types.py
@@ -223,7 +223,8 @@ class StructDef(RootType):
         if name == "__setstate__" or name not in self.fields:
             raise AttributeError(name)
         else:
-            return FieldReference(name, self)
+            return self.fields[name]
+            # return FieldReference(name, self)
 
     def accept(self, visitor: TypeVisitor[T]) -> T:
         return visitor.visit_struct_def(self)

I'm not sure why this wasn't implemented this way in the first place, but it does give us:

wat(Bakery.Packet.payload.v(V2.Message))

I'll have to dig more on why it wasn't done that way, not clear from comment. It definitely breaks compilation to change it though, so other compiler code wants the FieldReference.

[johnlevidy]

+    def __getitem__(self, match_struct_def: StructDef) -> Optional[T]:
+        for h in self.hash_types:
+            if isinstance(h, type(match_struct_def)):
+                return h
+        return None

better yet:

Bakery.Packet.payload[V2.Message]

[johnlevidy]

-        else:
+        elif isinstance(self.fields[name], EnumDef):
             return FieldReference(name, self)
+        else:
+            return self.fields[name]

seems to generate identical code. it seems to be a special case for enums, which we never have to reference. bit jank to have two behaviors? i think we need to do more surgery to really fix the enum case before we can make this work.

[johnlevidy]

Minor issues with above resolved:

wat(Bakery.Packet.payload[V2.Message].msg[V2.NewOrderResponse])

works with

index 20743d5..3151323 100644
--- a/python/tako/core/types.py
+++ b/python/tako/core/types.py
@@ -222,8 +222,10 @@ class StructDef(RootType):
         # Prevents infinite recursion on deepcopy
         if name == "__setstate__" or name not in self.fields:
             raise AttributeError(name)
-        else:
+        elif isinstance(self.fields[name], EnumDef):
             return FieldReference(name, self)
+        else:
+            return self.fields[name]
 
     def accept(self, visitor: TypeVisitor[T]) -> T:
         return visitor.visit_struct_def(self)
@@ -244,6 +246,11 @@ class VariantDef(RootType):
     tag_type: Int
     variants: t.Dict[StructDef, int]
 
+    def __getitem__(self, match_struct_def: StructDef) -> Optional[T]:
+        if match_struct_def in self.variants:
+            return match_struct_def
+        return None
+
     def accept(self, visitor: TypeVisitor[T]) -> T:
         return visitor.visit_variant_def(self)
 
@@ -269,6 +276,14 @@ class HashVariantDef(RootType):
     tag_type: Int
     hash_types: t.List[StructDef]
 
+    def __getitem__(self, match_struct_def: StructDef) -> Optional[T]:
+        for h in self.hash_types:
+            # TODO: is this totally safe? same name + all same fields.
+            # probably need a proper struct definition equality check?
+            if match_struct_def is h:
+                return h
+        return None
+
     def accept(self, visitor: TypeVisitor[T]) -> T:
         return visitor.visit_hash_variant_def(self)