Merge pull request #143 from pytest-dev/inject_into_class_ns

Allow factories to be registered within classes without polluting outer namespace.

CHANGES.rst

@@ -7,6 +7,7 @@ Unreleased
 - Improve "debuggability". Internal pytest-factoryboy calls are now visible when using a debugger like PDB or PyCharm.
 - Add type annotations. Now `register` and `LazyFixture` are type annotated.
 - Fix `_after_postgeneration` not getting the evaluated post_generations and RelatedFactory results correctly in the `result` param.
+- Factories can now be registered inside classes (even nested classes) and they won't pollute the module namespace.
 
 
 2.1.0

pytest_factoryboy/fixture.py

@@ -3,7 +3,7 @@
 
 import sys
 from dataclasses import dataclass
-from inspect import getmodule, signature
+from inspect import signature
 
 import factory
 import factory.builder
@@ -21,7 +21,6 @@
     from factory.builder import BuildStep
     from factory.declarations import PostGeneration
     from factory.declarations import PostGenerationContext
-    from types import ModuleType
 
     FactoryType = type[factory.Factory]
     T = TypeVar("T")
@@ -54,7 +53,6 @@ def register(factory_class: F, _name: str | None = None, **kwargs: Any) -> F:
 
     fixture_defs: list[FixtureDef] = []
 
-    module = get_caller_module()
     model_name = get_model_name(factory_class) if _name is None else _name
     factory_name = get_factory_name(factory_class)
 
@@ -117,7 +115,9 @@ def register(factory_class: F, _name: str | None = None, **kwargs: Any) -> F:
                     )
                 )
 
-    if not hasattr(module, factory_name):
+    caller_locals = get_caller_locals()
+
+    if factory_name not in caller_locals:
         fixture_defs.append(
             FixtureDef(
                 name=factory_name,
@@ -140,11 +140,31 @@ def register(factory_class: F, _name: str | None = None, **kwargs: Any) -> F:
 
     for fixture_def in fixture_defs:
         exported_name = fixture_def.name
-        setattr(module, exported_name, getattr(generated_module, exported_name))
+        fixture_function = getattr(generated_module, exported_name)
+        inject_into_caller(exported_name, fixture_function, caller_locals)
 
     return factory_class
 
 
+def inject_into_caller(name: str, function: Callable, locals_: dict[str, Any]) -> None:
+    """Inject a function into the caller's locals, making sure that the function will work also within classes."""
+    # We need to check if the caller frame is a class, since in that case the first argument is the class itself.
+    # In that case, we can apply the staticmethod() decorator to the injected function, so that the first param
+    # will be disregarded.
+    # To figure out if the caller frame is a class, we can check if the __qualname__ attribute is present.
+
+    # According to the python docs, __qualname__ is available for both **classes and functions**.
+    # However, it seems that for functions it is not yet available in the function namespace before it's defined.
+    # This could change in the future, but it shouldn't be too much of a problem since registering a factory
+    # in a function namespace would not make it usable anyway.
+    # Therefore, we can just check for __qualname__ to figure out if we are in a class, and apply the @staticmethod.
+    is_class_or_function = "__qualname__" in locals_
+    if is_class_or_function:
+        function = staticmethod(function)
+
+    locals_[name] = function
+
+
 def get_model_name(factory_class: FactoryType) -> str:
     """Get model fixture name by factory."""
     return (
@@ -338,14 +358,9 @@ def subfactory_fixture(request: FixtureRequest, factory_class: FactoryType) -> A
     return request.getfixturevalue(fixture)
 
 
-def get_caller_module(depth: int = 2) -> ModuleType:
-    """Get the module of the caller."""
-    frame = sys._getframe(depth)
-    module = getmodule(frame)
-    # Happens when there's no __init__.py in the folder
-    if module is None:
-        return get_caller_module(depth=depth)  # pragma: no cover
-    return module
+def get_caller_locals(depth: int = 2) -> dict[str, Any]:
+    """Get the local namespace of the caller frame."""
+    return sys._getframe(depth).f_locals
 
 
 class LazyFixture:

tests/test_namespaces.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass
+
+import pytest
+from _pytest.fixtures import FixtureLookupError
+from factory import Factory
+
+from pytest_factoryboy import register
+
+
+@dataclass
+class Foo:
+    value: str
+
+
+@register
+class FooFactory(Factory):
+    class Meta:
+        model = Foo
+
+    value = "module_foo"
+
+
+def test_module_namespace(foo):
+    assert foo.value == "module_foo"
+
+
+class TestClassNamespace:
+    @register
+    class FooFactory(Factory):
+        class Meta:
+            model = Foo
+
+        value = "class_foo"
+
+    register(FooFactory, "class_foo")
+
+    def test_class_namespace(self, class_foo, foo):
+        assert foo.value == class_foo.value == "class_foo"
+
+    class TestNestedClassNamespace:
+        @register
+        class FooFactory(Factory):
+            class Meta:
+                model = Foo
+
+            value = "nested_class_foo"
+
+        register(FooFactory, "nested_class_foo")
+
+        def test_nested_class_namespace(self, foo, nested_class_foo):
+            assert foo.value == nested_class_foo.value == "nested_class_foo"
+
+    def test_nested_class_factories_dont_pollute_the_class(self, request):
+        with pytest.raises(FixtureLookupError):
+            request.getfixturevalue("nested_class_foo")
+
+
+def test_class_factories_dont_pollute_the_module(request):
+    with pytest.raises(FixtureLookupError):
+        request.getfixturevalue("class_foo")
+    with pytest.raises(FixtureLookupError):
+        request.getfixturevalue("nested_class_foo")