From 4248f6390ec472fbdf407b235c824dae3c0dc06a Mon Sep 17 00:00:00 2001
From: Joe Taylor <joetaylorconsulting@gmail.com>
Date: Thu, 14 Jan 2021 22:27:37 +0100
Subject: [PATCH] Allow subclasses of UserObjects to be valid input even if
 they are defined after the object that references them (#53)

* Allow subclasses of UserObjects to be valid input even if they are defined after the object that references them

* Refactor to remove redundant code paths
---
 features/examples/modules/subclass_module.py | 14 +++---
 wysdom/__version__.py                        |  2 +-
 wysdom/object_schema/SchemaAnyOf.py          |  8 +++-
 wysdom/user_objects/UserObject.py            | 47 +++++++++++++++-----
 4 files changed, 50 insertions(+), 21 deletions(-)

diff --git a/features/examples/modules/subclass_module.py b/features/examples/modules/subclass_module.py
index 1dfa734..ebaa98e 100644
--- a/features/examples/modules/subclass_module.py
+++ b/features/examples/modules/subclass_module.py
@@ -1,4 +1,4 @@
-from typing import Dict, List
+from typing import List
 
 from abc import ABC, abstractmethod
 
@@ -28,14 +28,14 @@ def speak(self):
         return f"{self.name}, the greyhound, says Woof!"
 
 
+class Person(UserObject):
+    first_name: str = UserProperty(str)
+    last_name: str = UserProperty(str)
+    pets: List[Pet] = UserProperty(SchemaArray(Pet))
+
+
 class Cat(Pet):
     pet_type: str = UserProperty(SchemaConst("cat"))
 
     def speak(self):
         return f"{self.name} says Miaow!"
-
-
-class Person(UserObject):
-    first_name: str = UserProperty(str)
-    last_name: str = UserProperty(str)
-    pets: List[Pet] = UserProperty(SchemaArray(Pet))
diff --git a/wysdom/__version__.py b/wysdom/__version__.py
index 2a0cf82..7b491b5 100644
--- a/wysdom/__version__.py
+++ b/wysdom/__version__.py
@@ -1,3 +1,3 @@
-VERSION = (0, 2, 1)
+VERSION = (0, 2, 2)
 
 __version__ = '.'.join(map(str, VERSION))
diff --git a/wysdom/object_schema/SchemaAnyOf.py b/wysdom/object_schema/SchemaAnyOf.py
index 0338363..959c20b 100644
--- a/wysdom/object_schema/SchemaAnyOf.py
+++ b/wysdom/object_schema/SchemaAnyOf.py
@@ -15,7 +15,7 @@ class SchemaAnyOf(Schema):
                             is referred to by other schemas.
     """
 
-    allowed_schemas: Tuple[Schema] = None
+    _allowed_schemas: Tuple[Schema] = None
     schema_ref_name: Optional[str] = None
 
     def __init__(
@@ -23,7 +23,7 @@ def __init__(
             allowed_schemas: Iterable[Schema],
             schema_ref_name: Optional[str] = None
     ) -> None:
-        self.allowed_schemas = tuple(allowed_schemas)
+        self._allowed_schemas = tuple(allowed_schemas)
         self.schema_ref_name = schema_ref_name
 
     def __call__(
@@ -47,6 +47,10 @@ def __call__(
             )
         return valid_schemas[0](value, dom_info)
 
+    @property
+    def allowed_schemas(self) -> Tuple[Schema]:
+        return self._allowed_schemas
+
     @property
     def referenced_schemas(self) -> Dict[str, Schema]:
         referenced_schemas = {}
diff --git a/wysdom/user_objects/UserObject.py b/wysdom/user_objects/UserObject.py
index ae5eca9..0780f7c 100644
--- a/wysdom/user_objects/UserObject.py
+++ b/wysdom/user_objects/UserObject.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Type, Iterator, Union, Mapping
+from typing import Any, Type, Iterator, Union, Mapping, List
 
 import inspect
 
@@ -57,6 +57,38 @@ def _schema_superclasses(self) -> Iterator[Type[UserObject]]:
                 yield superclass
 
 
+class SchemaAnyRegisteredSubclass(SchemaAnyOf):
+    """
+    A SchemaAnyOf that allows any registered subclass of a UserObject.
+    Exists so that allowed_schemas can be populated dynamically and is not fixed
+    based on the object's current subclasses, but will also include any
+    future subclasses that are defined after this object's creation.
+
+    :param object_type:   The UserObject subclass.
+    """
+
+    def __init__(
+            self,
+            object_type: Type[UserObject]
+    ):
+        super().__init__(
+            allowed_schemas=[],
+            schema_ref_name=f"{object_type.__module__}.{object_type.__name__}"
+        )
+        assert issubclass(object_type, RegistersSubclasses)
+        self.object_type = object_type
+
+    @property
+    def allowed_schemas(self) -> List[Schema]:
+        return [
+            subclass.__json_schema__()
+            for subclass_list in self.object_type.registered_subclasses().values()
+            for subclass in subclass_list
+            if issubclass(subclass, UserObject)
+            and not isinstance(subclass.__json_schema__(), SchemaAnyOf)
+        ]
+
+
 class UserObject(DOMObject):
     """
     Base class for user-defined DOM objects.
@@ -110,16 +142,7 @@ def __json_schema__(cls) -> Schema:
             if cls.registered_subclasses():
                 has_subclasses = True
         if has_subclasses:
-            return SchemaAnyOf(
-                (
-                    subclass.__json_schema__()
-                    for subclass_list in cls.registered_subclasses().values()
-                    for subclass in subclass_list
-                    if issubclass(subclass, UserObject)
-                    and not isinstance(subclass.__json_schema__(), SchemaAnyOf)
-                ),
-                schema_ref_name=f"{cls.__module__}.{cls.__name__}"
-            )
+            return SchemaAnyRegisteredSubclass(cls)
         else:
             return SchemaObject(
                 properties=cls.__json_schema_properties__.properties,
@@ -128,3 +151,5 @@ def __json_schema__(cls) -> Schema:
                 object_type=cls,
                 schema_ref_name=f"{cls.__module__}.{cls.__name__}"
             )
+
+