Make EventQuerySet Generic (#170)

AmbitionEng · Oct 1, 2024 · a035c96 · a035c96
1 parent 87f1c05
commit a035c96
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/pghistory/models.py b/pghistory/models.py
@@ -1,5 +1,6 @@
 import uuid
 import warnings
+from typing import TYPE_CHECKING, TypeVar
 
 import django
 from django.apps import apps
@@ -11,6 +12,8 @@
 
 from pghistory import core, utils
 
+_M = TypeVar("_M", bound=models.Model)
+
 # This class is to preserve backwards compatibility with migrations
 PGHistoryJSONField = utils.JSONField
 
@@ -129,7 +132,13 @@ def get_compiler(self, *args, **kwargs):
         return compiler
 
 
-class EventQuerySet(models.QuerySet):
+if TYPE_CHECKING:
+    _EventQuerySetBase = models.QuerySet[_M]
+else:
+    _EventQuerySetBase = models.QuerySet
+
+
+class EventQuerySet(_EventQuerySetBase):
     """QuerySet with support for proxy fields"""
 
     def __init__(self, model=None, query=None, using=None, hints=None):