Merge pull request #22 from datacamp/fix/get_cls

[LE-1175] Fix get_cls in BaseNodeRegistry

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 All notable changes to this project will be documented in this file.
 
+## v0.8.1
+
+- Fix get_cls in BaseNodeRegistry, now updates fields of classes already in the registry
 
 ## v0.8.0
 

antlr_ast/__init__.py

@@ -1,3 +1,3 @@
-__version__ = "0.8.0"
+__version__ = "0.8.1"
 
 from . import ast

antlr_ast/ast.py

@@ -487,6 +487,12 @@ def get_cls(self, cls_name: str, field_names: tuple) -> Type[BaseNode]:
             self.dynamic_node_classes[cls_name] = BaseNode.create_cls(
                 cls_name, field_names
             )
+        else:
+            existing_cls = self.dynamic_node_classes[cls_name]
+            all_fields = tuple(set(existing_cls._fields) | set(field_names))
+            if len(all_fields) > len(existing_cls._fields):
+                existing_cls._fields = all_fields
+
         return self.dynamic_node_classes[cls_name]
 
     def isinstance(self, instance: BaseNode, class_name: str) -> bool:

tests/test_base_node_registry.py

@@ -0,0 +1,46 @@
+import pytest
+
+from antlr_ast.ast import BaseNodeRegistry
+
+
+def test_base_node_registry_get_cls():
+    # Given
+    base_node_registry = BaseNodeRegistry()
+
+    # When
+    cls1 = base_node_registry.get_cls("cls1", ("field1",))
+    cls1_2 = base_node_registry.get_cls("cls1", ("field1", "field2"))
+
+    # Then
+    assert set(cls1._fields) == {"field1", "field2"}
+    assert set(cls1_2._fields) == {"field1", "field2"}
+
+
+def test_base_node_registry_isinstance():
+    # Given
+    base_node_registry = BaseNodeRegistry()
+
+    # When
+    Cls1 = base_node_registry.get_cls("cls1", ("field1",))
+    Cls1_2 = base_node_registry.get_cls("cls1", ("field1", "field2"))
+    Cls2 = base_node_registry.get_cls("cls2", ("field_a", "field_b"))
+
+    cls1_obj = Cls1([], [], [])
+    cls1_2_obj = Cls1_2([], [], [])
+    cls2_obj = Cls2([], [], [])
+
+    # Then
+    assert isinstance(cls1_obj, type(cls1_2_obj))
+    assert base_node_registry.isinstance(cls1_obj, "cls1")
+    assert base_node_registry.isinstance(cls1_2_obj, "cls1")
+    assert base_node_registry.isinstance(cls2_obj, "cls2")
+    assert not base_node_registry.isinstance(cls1_obj, "cls2")
+    assert not base_node_registry.isinstance(cls1_2_obj, "cls2")
+    assert not base_node_registry.isinstance(cls2_obj, "cls1")
+
+    assert not base_node_registry.isinstance(cls2_obj, "cls3")
+
+    with pytest.raises(
+        TypeError, match="This function can only be used for BaseNode objects"
+    ):
+        base_node_registry.isinstance([], "cls1")