field_of function

Author: alessiostalla

CHANGELOG.md

@@ -2,10 +2,11 @@
 All notable changes to this project from version 0.4.0 upwards are documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
-## [0.9.0] – Not yet released
+## [0.9.0] – 2025-07-23
 
 ### Added
 - `PylasuANTLRParser` class modeled after the Kolasu equivalent
+- _field_of_ function to make AST transformers slightly more type-safe
 
 ## [0.8.1] – 2025-02-21
 

pylasu/transformation/transformation.py

@@ -1,5 +1,5 @@
 import functools
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, Field, fields
 from inspect import signature
 from typing import Any, Dict, Callable, TypeVar, Generic, Optional, List, Set, Iterable, Type, Union
 
@@ -29,6 +29,14 @@ def set(self, node: Node, value):
         return setattr(node, self.name, value)
 
 
+def field_of(cl: type, name: str) -> Field:
+    class_fields = fields(cl)
+    for fld in class_fields:
+        if fld.name == name:
+            return fld
+    raise Exception(f"Field {name} not found in {cl}")
+
+
 @dataclass
 class NodeFactory(Generic[Source, Output]):
     constructor: node_factory_constructor_type
@@ -37,8 +45,8 @@ class NodeFactory(Generic[Source, Output]):
 
     def with_child(
             self,
-            setter: Union[Callable[[Target, Optional[Child]], None], PropertyRef],
-            getter: Union[Callable[[Source], Optional[Any]], PropertyRef],
+            setter: Union[Callable[[Target, Optional[Child]], None], PropertyRef, Field],
+            getter: Union[Callable[[Source], Optional[Any]], PropertyRef, Field],
             name: Optional[str] = None,
             target_type: Optional[type] = None
     ) -> "NodeFactory[Source, Output]":
@@ -50,8 +58,12 @@ def with_child(
             prefix = ""
         if isinstance(getter, PropertyRef):
             getter = getter.get
+        elif isinstance(getter, Field):
+            getter = PropertyRef(getter.name).get
         if isinstance(setter, PropertyRef):
             setter = setter.set
+        elif isinstance(setter, Field):
+            setter = PropertyRef(setter.name).set
         self.children[prefix + name] = ChildNodeFactory(prefix + name, getter, setter)
         return self
 

tests/mapping/test_parse_tree_to_ast_transformers.py

@@ -6,7 +6,7 @@
 
 from pylasu.mapping.parse_tree_to_ast_transformer import ParseTreeToASTTransformer
 from pylasu.model import Node, Named, ReferenceByName
-from pylasu.transformation.transformation import PropertyRef
+from pylasu.transformation.transformation import PropertyRef, field_of
 from tests.antlr_entity.AntlrEntityLexer import AntlrEntityLexer
 from tests.antlr_entity.AntlrEntityParser import AntlrEntityParser
 
@@ -50,7 +50,7 @@ class ParseTreeToASTTransformerTest(unittest.TestCase):
     def test_simple_entities_transformer(self):
         transformer = ParseTreeToASTTransformer(allow_generic_node=False)
         transformer.register_node_factory(AntlrEntityParser.ModuleContext, lambda ctx: EModule(name=ctx.name.text)) \
-            .with_child(PropertyRef("entities"), AntlrEntityParser.ModuleContext.entity)
+            .with_child(field_of(EModule, "entities"), AntlrEntityParser.ModuleContext.entity)
         transformer.register_node_factory(AntlrEntityParser.EntityContext, lambda ctx: EEntity(name=ctx.name.text))
         expected_ast = EModule("M", [EEntity("FOO", []), EEntity("BAR", [])])
         actual_ast = transformer.transform(self.parse_entities("""