Add support for SQLAlchemyFactory and generic ModelFactory types

Author: vrslev

src/community_of_python_flake8_plugin/checks/dataclass_config.py

@@ -54,9 +54,9 @@ def is_pydantic_model(class_node: ast.ClassDef) -> bool:
 def is_model_factory(class_node: ast.ClassDef) -> bool:
     """Check if class inherits from ModelFactory."""
     for base_class in class_node.bases:
-        if isinstance(base_class, ast.Name) and base_class.id == "ModelFactory":
+        if isinstance(base_class, ast.Name) and base_class.id in {"ModelFactory", "SQLAlchemyFactory"}:
             return True
-        if isinstance(base_class, ast.Attribute) and base_class.attr == "ModelFactory":
+        if isinstance(base_class, ast.Attribute) and base_class.attr in {"ModelFactory", "SQLAlchemyFactory"}:
             return True
     return False
 

src/community_of_python_flake8_plugin/constants.py

@@ -108,6 +108,7 @@
     "ensure",
     "submit",
     "clear",
+    "undo",
 }
 
 SCALAR_ANNOTATIONS: typing.Final = {"int", "str", "float", "bool", "bytes", "complex"}
@@ -116,5 +117,5 @@
 
 ALLOWED_STDLIB_FROM_IMPORTS: typing.Final = {"collections.abc"}
 
-FINAL_CLASS_EXCLUDED_BASES: typing.Final = {"BaseModel", "RootModel", "ModelFactory"}
+FINAL_CLASS_EXCLUDED_BASES: typing.Final = {"BaseModel", "RootModel", "ModelFactory", "SQLAlchemyFactory"}
 MAX_IMPORT_NAMES: typing.Final = 2

src/community_of_python_flake8_plugin/utils.py

@@ -17,6 +17,12 @@ def check_inherits_from_bases(class_definition: ast.ClassDef, base_classes: set[
             return True
         if isinstance(base_class, ast.Attribute) and base_class.attr in base_classes:
             return True
+        # Handle generic types like ModelFactory[SomeType]
+        if isinstance(base_class, ast.Subscript):
+            if isinstance(base_class.value, ast.Name) and base_class.value.id in base_classes:
+                return True
+            if isinstance(base_class.value, ast.Attribute) and base_class.value.attr in base_classes:
+                return True
     return False
 
 

tests/test_plugin.py

@@ -187,6 +187,24 @@ def test_type_annotation_validations(input_source: str, expected_output: list[st
             "from polyfactory.factories.pydantic_factory import ModelFactory\nclass MyFactory(ModelFactory):\n    def calculator(self): pass",
             ["COP012"],
         ),
+        # No violation: ModelFactory generic methods should be exempt from COP009
+        (
+            "from polyfactory.factories.pydantic_factory import ModelFactory\nimport some_module\n"
+            "class MyFactory(ModelFactory[some_module.SomeClass]):\n    def calculator(self): pass",
+            ["COP012"],
+        ),
+        # No violation: ModelFactory classmethod should be exempt from COP009
+        (
+            "from polyfactory.factories.pydantic_factory import ModelFactory\n"
+            "class MyFactory(ModelFactory):\n    @classmethod\n    def create(cls): pass",
+            ["COP012"],
+        ),
+        # No violation: ModelFactory generic classmethod should be exempt from COP009
+        (
+            "from polyfactory.factories.pydantic_factory import ModelFactory\nimport some_module\n"
+            "class MyFactory(ModelFactory[some_module.SomeClass]):\n    @classmethod\n    def create(cls): pass",
+            ["COP012"],
+        ),
         # No violation: cached_property imported directly should exempt function from COP009 (but triggers COP002 for import style)
         (
             "from functools import cached_property\nclass ExampleClass:\n    @cached_property\n    def calculator(self): pass",