Add support for finding model classes in modules (#58)

* Add function to find models in a module * Resolve generator in test * Create, draw, to_dot module support * Support passing modules in CLI * Update changelog * Add issue and PR links to changelog * Update changelog Co-authored-by: Jay Qi <jayqi@users.noreply.github.com>
drivendataorg · Jul 30, 2022 · adab9ba · adab9ba
1 parent df06f7e
commit adab9ba
Show file tree

Hide file tree

Showing 8 changed files with 393 additions and 40 deletions.
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,8 +1,11 @@
 # erdantic Changelog
 
-## v0.5.0 (2022-06-03)
+## v0.5.0 (2022-07-29)
 
 - Removed support for Python 3.6. ([Issue #51](https://github.com/drivendataorg/erdantic/issues/51), [PR #56](https://github.com/drivendataorg/erdantic/pull/56))
+- Added support for modules as inputs to all entrypoints to diagram creation (`create`, `draw`, `to_dot`, CLI). For all modules passed, erdantic will find all supported data model classes in each module. ([Issue #23](https://github.com/drivendataorg/erdantic/issues/23), [PR #58](https://github.com/drivendataorg/erdantic/pull/58))
+    - Added new parameter `limit_search_models_to` to all entrypoints to allow for limiting which data model classes will be yielded from searching a module.
+
 
 ## v0.4.1 (2022-04-08)
 

diff --git a/erdantic/base.py b/erdantic/base.py
@@ -2,7 +2,7 @@
 import inspect
 from typing import Any, Callable, Dict, Generic, List, Optional, Type, TypeVar, Union
 
-from erdantic.exceptions import InvalidModelAdapterError
+from erdantic.exceptions import InvalidModelAdapterError, ModelAdapterNotFoundError
 from erdantic.typing import Final, GenericAlias, repr_type
 
 
@@ -207,3 +207,17 @@ def decorator(cls: type) -> type:
         return cls
 
     return decorator
+
+
+def get_model_adapter(key_or_adapter: Union[str, type]):
+    if isinstance(key_or_adapter, str):
+        try:
+            return model_adapter_registry[key_or_adapter]
+        except KeyError:
+            raise ModelAdapterNotFoundError(
+                f"No model adapter registered with key {key_or_adapter}"
+            )
+    elif isinstance(key_or_adapter, type) and issubclass(key_or_adapter, Model):
+        return key_or_adapter
+    else:
+        raise InvalidModelAdapterError("Input must be str or subclass of erdantic.base.Model.")
diff --git a/erdantic/cli.py b/erdantic/cli.py
@@ -1,16 +1,36 @@
+from enum import Enum
 from importlib import import_module
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Optional, TYPE_CHECKING
 
 import typer
 
+from erdantic.base import model_adapter_registry
 from erdantic.erd import create
+from erdantic.exceptions import ModelOrModuleNotFoundError
 from erdantic.version import __version__
 
 
 app = typer.Typer()
 
 
+class StrEnum(str, Enum):
+    pass
+
+
+if TYPE_CHECKING:
+    # mypy typechecking doesn't really support enums created with functional API
+    # https://github.com/python/mypy/issues/6037
+
+    class SupportedModelIdentifier(StrEnum):
+        pass
+
+else:
+    SupportedModelIdentifier = StrEnum(
+        "SupportedModelIdentifier", {key: key for key in model_adapter_registry.keys()}
+    )
+
+
 def version_callback(version: bool):
     """Print erdantic version to console."""
     if version:
@@ -29,12 +49,13 @@ def dot_callback(ctx: typer.Context, dot: bool):
 
 @app.command()
 def main(
-    models: List[str] = typer.Argument(
+    models_or_modules: List[str] = typer.Argument(
         ...,
         help=(
-            "One or more full dotted paths for data model classes to include in diagram, "
-            "e.g., 'erdantic.examples.pydantic.Party'. Only the root models of composition trees "
-            "are needed; erdantic will traverse the composition tree to find component classes."
+            "One or more full dotted paths for data model classes, or modules containing data "
+            "model classes, to include in diagram, e.g., 'erdantic.examples.pydantic.Party'. Only "
+            "the root models of composition trees are needed; erdantic will traverse the "
+            "composition tree to find component classes."
         ),
     ),
     termini: List[str] = typer.Option(
@@ -47,6 +68,17 @@ def main(
             "Repeat this option if more than one."
         ),
     ),
+    limit_search_models_to: List[SupportedModelIdentifier] = typer.Option(
+        None,
+        "--limit-search-models-to",
+        "-m",
+        help=(
+            "Identifiers of model classes that erdantic supports. If any are specified, when "
+            "searching a module, limit data model classes to those ones. Repeat this option if "
+            "more than one.Defaults to None which will find all data model classes supported by "
+            "erdantic. "
+        ),
+    ),
     out: Path = typer.Option(..., "--out", "-o", help="Output filename."),
     dot: Optional[bool] = typer.Option(
         None,
@@ -73,18 +105,16 @@ def main(
     rendered using the Graphviz library. Currently supported data modeling frameworks are Pydantic
     and standard library dataclasses.
     """
-    model_classes: List[type] = []
-    for model in models:
-        module_name, model_name = model.rsplit(".", 1)
-        module = import_module(module_name)
-        model_classes.append(getattr(module, model_name))
-    termini_classes: List[type] = []
-    for terminus in termini:
-        module_name, model_name = terminus.rsplit(".", 1)
-        module = import_module(module_name)
-        termini_classes.append(getattr(module, model_name))
-
-    diagram = create(*model_classes, termini=termini_classes)
+    model_or_module_objs = [import_object_from_name(mm) for mm in models_or_modules]
+    termini_classes = [import_object_from_name(mm) for mm in termini]
+    limit_search_models_to_str = [
+        m.value for m in limit_search_models_to
+    ] or None  # Don't want empty list
+    diagram = create(
+        *model_or_module_objs,
+        termini=termini_classes,
+        limit_search_models_to=limit_search_models_to_str,
+    )
     if dot:
         typer.echo(diagram.to_dot())
     else:
@@ -93,3 +123,16 @@ def main(
             raise typer.Exit(code=1)
         diagram.draw(out)
         typer.echo(f"Rendered diagram to {out}")
+
+
+def import_object_from_name(full_obj_name):
+    # Try to import as a module
+    try:
+        return import_module(full_obj_name)
+    except ModuleNotFoundError:
+        try:
+            module_name, obj_name = full_obj_name.rsplit(".", 1)
+            module = import_module(module_name)
+            return getattr(module, obj_name)
+        except (ImportError, AttributeError):
+            raise ModelOrModuleNotFoundError(f"{full_obj_name} not found")
diff --git a/erdantic/erd.py b/erdantic/erd.py
@@ -1,9 +1,11 @@
+import inspect
 import os
-from typing import Any, List, Sequence, Set, Union
+from types import ModuleType
+from typing import Any, Iterable, Iterator, List, Optional, Sequence, Set, Type, Union
 
 import pygraphviz as pgv
 
-from erdantic.base import Field, Model, model_adapter_registry
+from erdantic.base import Field, get_model_adapter, Model, model_adapter_registry
 from erdantic.exceptions import (
     NotATypeError,
     _StringForwardRefError,
@@ -166,24 +168,39 @@ def _repr_svg_(self) -> str:
         return graph.draw(prog="dot", format="svg").decode(graph.encoding)
 
 
-def create(*models: type, termini: Sequence[type] = []) -> EntityRelationshipDiagram:
+def create(
+    *models_or_modules: Union[type, ModuleType],
+    termini: Sequence[type] = [],
+    limit_search_models_to: Optional[Iterable[str]] = None,
+) -> EntityRelationshipDiagram:
     """Construct [`EntityRelationshipDiagram`][erdantic.erd.EntityRelationshipDiagram] from given
     data model classes.
 
     Args:
-        *models (type): Data model classes to diagram.
+        *models_or_modules (type): Data model classes to diagram or modules containing them.
         termini (Sequence[type]): Data model classes to set as terminal nodes. erdantic will stop
             searching for component classes when it reaches these models
-
+        limit_search_models_to (Optional[Iterable[sr]], optional): Iterable of identifiers of data
+            model classes that erdantic supports. If any are specified, when searching a module,
+            limit data model classes to those ones. Defaults to None which will find all data model
+            classes supported by erdantic.
     Raises:
         UnknownModelTypeError: If model is not recognized as a supported model type.
 
     Returns:
         EntityRelationshipDiagram: diagram object for given data model.
     """
-    for raw_model in models + tuple(termini):
-        if not isinstance(raw_model, type):
-            raise NotATypeError(f"Given model is not a type: {raw_model}")
+    models = []
+    for mm in models_or_modules:
+        if isinstance(mm, type):
+            models.append(mm)
+        elif isinstance(mm, ModuleType):
+            models.extend(find_models(mm, limit_search_models_to=limit_search_models_to))
+        else:
+            raise NotATypeError(f"Given model is not a type: {mm}")
+    for terminal_model in tuple(termini):
+        if not isinstance(terminal_model, type):
+            raise NotATypeError(f"Given terminal model is not a type: {terminal_model}")
 
     seen_models: Set[Model] = {adapt_model(t) for t in termini}
     seen_edges: Set[Edge] = set()
@@ -193,6 +210,34 @@ def create(*models: type, termini: Sequence[type] = []) -> EntityRelationshipDia
     return EntityRelationshipDiagram(models=list(seen_models), edges=list(seen_edges))
 
 
+def find_models(
+    module: ModuleType, limit_search_models_to: Optional[Iterable[str]] = None
+) -> Iterator[type]:
+    """Searches a module and yields all data model classes found.
+
+    Args:
+        module (ModuleType): Module to search for data model classes
+        limit_search_models_to (Optional[Iterable[sr]], optional): Iterable of identifiers of data
+            model classes that erdantic supports. If any are specified, when searching a module,
+            limit data model classes to those ones. Defaults to None which will find all data model
+            classes supported by erdantic.
+
+    Yields:
+        Iterator[type]: Members of module that are data model classes.
+    """
+    limit_search_models_to_adapters: Iterable[Type[Model]]
+    if limit_search_models_to is None:
+        limit_search_models_to_adapters = model_adapter_registry.values()
+    else:
+        limit_search_models_to_adapters = [get_model_adapter(m) for m in limit_search_models_to]
+
+    for _, member in inspect.getmembers(module, inspect.isclass):
+        if member.__module__ == module.__name__:
+            for model_adapter in limit_search_models_to_adapters:
+                if model_adapter.is_model_type(member):
+                    yield member
+
+
 def adapt_model(obj: Any) -> Model:
     """Dispatch object to appropriate concrete [`Model`][erdantic.base.Model] adapter subclass and
     return instantiated adapter instance.
@@ -247,31 +292,53 @@ def search_composition_graph(
                 ) from None
 
 
-def draw(*models: type, out: Union[str, os.PathLike], termini: Sequence[type] = [], **kwargs):
+def draw(
+    *models_or_modules: Union[type, ModuleType],
+    out: Union[str, os.PathLike],
+    termini: Sequence[type] = [],
+    limit_search_models_to: Optional[Iterable[str]] = None,
+    **kwargs,
+):
     """Render entity relationship diagram for given data model classes to file.
 
     Args:
-        *models (type): Data model classes to diagram.
+        *models_or_modules (type): Data model classes to diagram, or modules containing them.
         out (Union[str, os.PathLike]): Output file path for rendered diagram.
         termini (Sequence[type]): Data model classes to set as terminal nodes. erdantic will stop
             searching for component classes when it reaches these models
+        limit_search_models_to (Optional[Iterable[sr]], optional): Iterable of identifiers of data
+            model classes that erdantic supports. If any are specified, when searching a module,
+            limit data model classes to those ones. Defaults to None which will find all data model
+            classes supported by erdantic.
         **kwargs: Additional keyword arguments to [`pygraphviz.AGraph.draw`](https://pygraphviz.github.io/documentation/latest/reference/agraph.html#pygraphviz.AGraph.draw).
     """
-    diagram = create(*models, termini=termini)
+    diagram = create(
+        *models_or_modules, termini=termini, limit_search_models_to=limit_search_models_to
+    )
     diagram.draw(out=out, **kwargs)
 
 
-def to_dot(*models: type, termini: Sequence[type] = []) -> str:
+def to_dot(
+    *models_or_modules: Union[type, ModuleType],
+    termini: Sequence[type] = [],
+    limit_search_models_to: Optional[Iterable[str]] = None,
+) -> str:
     """Generate Graphviz [DOT language](https://graphviz.org/doc/info/lang.html) representation of
     entity relationship diagram for given data model classes.
 
     Args:
-        *models (type): Data model classes to diagram.
+        *models_or_modules (type): Data model classes to diagram, or modules containing them.
         termini (Sequence[type]): Data model classes to set as terminal nodes. erdantic will stop
             searching for component classes when it reaches these models
+        limit_search_models_to (Optional[Iterable[sr]], optional): Iterable of identifiers of data
+            model classes that erdantic supports. If any are specified, when searching a module,
+            limit data model classes to those ones. Defaults to None which will find all data model
+            classes supported by erdantic.
 
     Returns:
         str: DOT language representation of diagram
     """
-    diagram = create(*models, termini=termini)
+    diagram = create(
+        *models_or_modules, termini=termini, limit_search_models_to=limit_search_models_to
+    )
     return diagram.to_dot()
diff --git a/erdantic/exceptions.py b/erdantic/exceptions.py
@@ -15,13 +15,21 @@ class InvalidModelError(ValueError, ErdanticException):
 
 
 class InvalidModelAdapterError(ValueError, ErdanticException):
-    """Raised when trying to register a model adapter that is not subclassing Model."""
+    """Raised when a model adapter is expected but input is not subclassing Model."""
 
 
 class InvalidFieldError(ValueError, ErdanticException):
     """Raised when an invalid field object is passed to a field adapter."""
 
 
+class ModelAdapterNotFoundError(KeyError, ErdanticException):
+    """Raised when specified key does not match a registered model adapter."""
+
+
+class ModelOrModuleNotFoundError(ImportError, ErdanticException):
+    """Raised when specified fully qualified name of model class or module cannot be found."""
+
+
 class NotATypeError(ValueError, ErdanticException):
     pass
 

diff --git a/tests/test_base.py b/tests/test_base.py
@@ -1,9 +1,10 @@
 import pytest
 
 import erdantic as erd
-from erdantic.base import Field, Model, register_model_adapter
+from erdantic.base import Field, get_model_adapter, Model, register_model_adapter
 from erdantic.examples.pydantic import Party
-from erdantic.exceptions import InvalidModelAdapterError
+from erdantic.exceptions import InvalidModelAdapterError, ModelAdapterNotFoundError
+from erdantic.pydantic import PydanticModel
 
 
 def test_abstract_field_instatiation():
@@ -31,3 +32,12 @@ def test_repr():
     diagram = erd.create(Party)
     assert repr(diagram.models[0]) and isinstance(repr(diagram.models[0]), str)
     assert repr(diagram.models[0].fields[0]) and isinstance(repr(diagram.models[0].fields[0]), str)
+
+
+def test_get_model_adapter():
+    assert get_model_adapter("pydantic") == PydanticModel
+    assert get_model_adapter(PydanticModel) == PydanticModel
+    with pytest.raises(ModelAdapterNotFoundError):
+        get_model_adapter("unknown_key")
+    with pytest.raises(InvalidModelAdapterError):
+        get_model_adapter(Party)