Fix error from type annotations containing ellipsis (#127)

* Fix error from type annotations containing ellipsis * Fix wrong number --------- Co-authored-by: Jay Qi <[email protected]>
drivendataorg · Sep 19, 2024 · 3464382 · 3464382
1 parent e77348e
commit 3464382
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 2 deletions.
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,5 +1,9 @@
 # erdantic Changelog
 
+## v1.0.5 (Unreleased)
+
+- Fixed runtime `AttributeError` that occurred when creating a diagram that includes a model with a field that uses a type annotation with ellipsis literal (e.g., `tuple[int, ...]`). ([Issue #124](https://github.com/drivendataorg/erdantic/issues/124), [PR #127](https://github.com/drivendataorg/erdantic/pull/127))
+
 ## v1.0.4 (2024-07-16)
 
 - Fixed handling of `typing.Annotated` in cases where it's not the outermost generic type. ([Issue #122](https://github.com/drivendataorg/erdantic/issues/122), [PR #123](https://github.com/drivendataorg/erdantic/pull/123))

diff --git a/erdantic/core.py b/erdantic/core.py
@@ -512,6 +512,10 @@ def _add_if_model(self, model: type, recurse: bool) -> bool:
             # These are not going to be models
             if "__qualname__" in str(e):
                 return False
+            # ellipsis object (used for example in tuple[int, ...]) don't have __module__ attribute
+            # This is also not going to be a model
+            elif "__module__" in str(e):
+                return False
             raise
         if key not in self.models:
             try:

diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "erdantic"
-version = "1.0.4"
+version = "1.0.5"
 description = "Entity relationship diagrams for Python data model classes like Pydantic."
 readme = "README.md"
 authors = [{ name = "DrivenData", email = "[email protected]" }, { name = "Jay Qi" }]

diff --git a/tests/test_core.py b/tests/test_core.py
@@ -4,7 +4,7 @@
 import os
 from pathlib import Path
 import sys
-from typing import Any, AnyStr, List, Literal, Optional, TypeVar
+from typing import Any, AnyStr, List, Literal, Optional, Tuple, TypeVar
 
 if sys.version_info >= (3, 9):
     from typing import Annotated
@@ -200,6 +200,18 @@ class MyModel(pydantic.BaseModel):
 
     diagram = EntityRelationshipDiagram()
     diagram.add_model(MyModel)
+    diagram.to_dot()
+
+
+def test_model_with_ellipsis():
+    """Model with Ellipsis should not error."""
+
+    class EllipsisModel(pydantic.BaseModel):
+        ellipsis_field: Tuple[int, ...]
+
+    diagram = EntityRelationshipDiagram()
+    diagram.add_model(EllipsisModel)
+    diagram.to_dot()
 
 
 def test_model_with_no_fields():