Add support for swapping implementations by subclassing

Support for OptimizationCompiler, FilterInfoCompiler, QueryOptimizer,
and FieldSelectionCompiler.

query_optimizer/compiler.py

@@ -13,7 +13,7 @@
 from .optimizer import QueryOptimizer
 from .prefetch_hack import fetch_in_context
 from .settings import optimizer_settings
-from .utils import is_optimized, optimizer_logger
+from .utils import is_optimized, optimizer_logger, swappable_by_subclassing
 
 if TYPE_CHECKING:
     import graphene
@@ -67,6 +67,7 @@ def optimize_single(
     return next(iter(queryset), None)
 
 
+@swappable_by_subclassing
 class OptimizationCompiler(GraphQLASTWalker):
     """Class for compiling SQL optimizations based on the given query."""
 

query_optimizer/filter_info.py

@@ -18,6 +18,7 @@
     is_node,
 )
 from .typing import GQLInfo, GraphQLFilterInfo, ToManyField, ToOneField
+from .utils import swappable_by_subclassing
 
 if TYPE_CHECKING:
     from django.db.models import Model
@@ -39,6 +40,7 @@ def get_filter_info(info: GQLInfo, model: type[Model]) -> GraphQLFilterInfo:
     return compiler.filter_info.get(to_snake_case(info.field_name), {})
 
 
+@swappable_by_subclassing
 class FilterInfoCompiler(GraphQLASTWalker):
     """Class for compiling filtering information from a GraphQL query."""
 

query_optimizer/optimizer.py

@@ -23,6 +23,7 @@
     calculate_slice_for_queryset,
     mark_optimized,
     optimizer_logger,
+    swappable_by_subclassing,
 )
 from .validators import validate_pagination_args
 
@@ -43,6 +44,7 @@ class CompilationResults:
     prefetch_related: list[Prefetch | str] = dataclasses.field(default_factory=list)
 
 
+@swappable_by_subclassing
 class QueryOptimizer:
     """Creates optimized queryset based on the optimization data found by the OptimizationCompiler."""
 

query_optimizer/selections.py

@@ -7,6 +7,8 @@
 
 from query_optimizer.ast import GraphQLASTWalker, get_selections
 
+from .utils import swappable_by_subclassing
+
 if TYPE_CHECKING:
     from django.db import models
     from graphene.types.definitions import GrapheneObjectType
@@ -27,6 +29,7 @@ def get_field_selections(info: GQLInfo, model: Optional[type[models.Model]] = No
     return compiler.field_selections[0][to_snake_case(info.field_name)]
 
 
+@swappable_by_subclassing
 class FieldSelectionCompiler(GraphQLASTWalker):
     """Class for compiling filtering information from a GraphQL query."""
 

query_optimizer/utils.py

@@ -8,10 +8,11 @@
 from .settings import optimizer_settings
 
 if TYPE_CHECKING:
-    from .typing import Optional, ParamSpec, TypeVar, Union
+    from .typing import Any, Optional, ParamSpec, TypeVar, Union
 
     T = TypeVar("T")
     P = ParamSpec("P")
+    Ttype = TypeVar("Ttype", bound=type)
 
 
 __all__ = [
@@ -22,6 +23,7 @@
     "mark_optimized",
     "optimizer_logger",
     "remove_optimized_mark",
+    "swappable_by_subclassing",
 ]
 
 
@@ -189,3 +191,23 @@ def add_slice_to_queryset(
 class SubqueryCount(models.Subquery):
     template = "(SELECT COUNT(*) FROM (%(subquery)s) _count)"
     output_field = models.BigIntegerField()
+
+
+def swappable_by_subclassing(obj: Ttype) -> Ttype:
+    """Makes the decorated class return the most recently created direct subclass when it is instantiated."""
+    orig_init_subclass = obj.__init_subclass__
+
+    def init_subclass(*args: Any, **kwargs: Any) -> None:
+        nonlocal obj
+
+        new_subcls: type = obj.__subclasses__()[-1]
+
+        def new(_: type, *args: Any, **kwargs: Any) -> Ttype:
+            return super(type, new_subcls).__new__(new_subcls, *args, **kwargs)
+
+        obj.__new__ = new
+
+        return orig_init_subclass(*args, **kwargs)
+
+    obj.__init_subclass__ = init_subclass
+    return obj

tests/test_utils.py

@@ -3,7 +3,7 @@
 
 from query_optimizer.settings import optimizer_settings
 from query_optimizer.typing import NamedTuple, Optional
-from query_optimizer.utils import calculate_queryset_slice, calculate_slice_for_queryset
+from query_optimizer.utils import calculate_queryset_slice, calculate_slice_for_queryset, swappable_by_subclassing
 from tests.example.models import Example
 from tests.factories.example import ExampleFactory
 from tests.helpers import parametrize_helper
@@ -144,3 +144,40 @@ def test_calculate_slice_for_queryset(pagination_input: PaginationInput, start:
     )
 
     assert values == {"start": start, "stop": stop}
+
+
+def test_swappable_by_subclassing():
+    @swappable_by_subclassing
+    class A:
+        def __init__(self) -> None:
+            self.one = 1
+
+    a = A()
+    assert type(a) is A
+    assert a.one == 1
+
+    class B(A):
+        def __init__(self) -> None:
+            super().__init__()
+            self.two = 2
+
+    b = A()
+    assert type(b) is B
+    assert b.one == 1
+    assert b.two == 2
+
+    class C(A):
+        def __init__(self) -> None:
+            super().__init__()
+            self.three = 3
+
+    c = A()
+    assert type(c) is C
+    assert c.one == 1
+    assert not hasattr(c, "two")
+    assert c.three == 3
+
+    class D(B): ...
+
+    d = A()
+    assert type(d) is C  # Only direct subclasses are swapped.