feat: add mixins for CommittableQuerySet and PatientAssetQuerySet (#211)

canvas_sdk/v1/data/allergy_intolerance.py

@@ -1,6 +1,26 @@
+from typing import cast
+
 from django.db import models
 
-from canvas_sdk.v1.data.base import CommittableModelManager, ValueSetLookupQuerySet
+from canvas_sdk.v1.data.base import (
+    BaseModelManager,
+    CommittableQuerySetMixin,
+    ForPatientQuerySetMixin,
+    ValueSetLookupQuerySet,
+)
+
+
+class AllergyIntoleranceQuerySet(
+    ValueSetLookupQuerySet,
+    CommittableQuerySetMixin,
+    ForPatientQuerySetMixin,
+):
+    """AllergyIntoleranceQuerySet."""
+
+    pass
+
+
+AllergyIntoleranceManager = BaseModelManager.from_queryset(AllergyIntoleranceQuerySet)
 
 
 class AllergyIntolerance(models.Model):
@@ -10,7 +30,7 @@ class Meta:
         managed = False
         db_table = "canvas_sdk_data_api_allergyintolerance_001"
 
-    objects = CommittableModelManager().from_queryset(ValueSetLookupQuerySet)()
+    objects = cast(AllergyIntoleranceQuerySet, AllergyIntoleranceManager())
 
     id = models.UUIDField()
     dbid = models.BigIntegerField(primary_key=True)

canvas_sdk/v1/data/base.py

@@ -10,25 +10,12 @@
     from canvas_sdk.value_set.value_set import ValueSet
 
 
-class CommittableModelManager(models.Manager):
-    """A manager for commands that can be committed."""
+class BaseModelManager(models.Manager):
+    """A base manager for models."""
 
-    def get_queryset(self) -> "CommittableQuerySet":
+    def get_queryset(self) -> models.QuerySet:
         """Return a queryset that filters out deleted objects."""
-        # TODO: Should we just filter these out at the view level?
-        return CommittableQuerySet(self.model, using=self._db).filter(deleted=False)
-
-
-class CommittableQuerySet(models.QuerySet):
-    """A queryset for committable objects."""
-
-    def committed(self) -> "Self":
-        """Return a queryset that filters for objects that have been committed."""
-        return self.filter(committer_id__isnull=False, entered_in_error_id__isnull=True)
-
-    def for_patient(self, patient_id: str) -> "Self":
-        """Return a queryset that filters objects for a specific patient."""
-        return self.filter(patient__id=patient_id)
+        return super().get_queryset().filter(deleted=False)
 
 
 class BaseQuerySet(models.QuerySet):
@@ -40,10 +27,14 @@ class BaseQuerySet(models.QuerySet):
 class QuerySetProtocol(Protocol):
     """A typing protocol for use in mixins into models.QuerySet-inherited classes."""
 
-    def filter(self, *args: Any, **kwargs: Any) -> models.QuerySet[Any]:
+    def filter(self, *args: Any, **kwargs: Any) -> Self:
         """Django's models.QuerySet filter method."""
         ...
 
+    def distinct(self) -> Self:
+        """Django's models.QuerySet distinct method."""
+        ...
+
 
 class ValueSetLookupQuerySetProtocol(QuerySetProtocol):
     """A typing protocol for use in mixins using value set lookup methods."""
@@ -61,10 +52,26 @@ def q_object(system: str, codes: Container[str]) -> Q:
         raise NotImplementedError
 
 
+class CommittableQuerySetMixin(QuerySetProtocol):
+    """A queryset for committable objects."""
+
+    def committed(self) -> Self:
+        """Return a queryset that filters for objects that have been committed."""
+        return self.filter(committer_id__isnull=False, entered_in_error_id__isnull=True)
+
+
+class ForPatientQuerySetMixin(QuerySetProtocol):
+    """A queryset for patient assets."""
+
+    def for_patient(self, patient_id: str) -> Self:
+        """Return a queryset that filters objects for a specific patient."""
+        return self.filter(patient__id=patient_id)
+
+
 class ValueSetLookupQuerySetMixin(ValueSetLookupQuerySetProtocol):
     """A QuerySet mixin that can filter objects based on a ValueSet."""
 
-    def find(self, value_set: type["ValueSet"]) -> models.QuerySet[Any]:
+    def find(self, value_set: type["ValueSet"]) -> Self:
         """
         Filters conditions, medications, etc. to those found in the inherited ValueSet class that is passed.
 
@@ -146,7 +153,7 @@ def timeframe_filter_field(self) -> str:
         """Returns the field that should be filtered on. Can be overridden for different models."""
         return "note__datetime_of_service"
 
-    def within(self, timeframe: "Timeframe") -> models.QuerySet:
+    def within(self, timeframe: "Timeframe") -> Self:
         """A method to filter a queryset for datetimes within a timeframe."""
         return self.filter(
             **{
@@ -158,13 +165,19 @@ def within(self, timeframe: "Timeframe") -> models.QuerySet:
         )
 
 
-class ValueSetLookupQuerySet(CommittableQuerySet, ValueSetLookupQuerySetMixin):
+class CommittableQuerySet(BaseQuerySet, CommittableQuerySetMixin):
+    """A queryset for committable objects."""
+
+    pass
+
+
+class ValueSetLookupQuerySet(BaseQuerySet, ValueSetLookupQuerySetMixin):
     """A class that includes methods for looking up value sets."""
 
     pass
 
 
-class ValueSetLookupByNameQuerySet(CommittableQuerySet, ValueSetLookupByNameQuerySetMixin):
+class ValueSetLookupByNameQuerySet(BaseQuerySet, ValueSetLookupByNameQuerySetMixin):
     """A class that includes methods for looking up value sets by name."""
 
     pass

canvas_sdk/v1/data/billing.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Self
 
 from django.db import models
 
@@ -12,7 +12,7 @@
 class BillingLineItemQuerySet(ValueSetTimeframeLookupQuerySet):
     """A class that adds functionality to filter BillingLineItem objects."""
 
-    def find(self, value_set: type["ValueSet"]) -> models.QuerySet:
+    def find(self, value_set: type["ValueSet"]) -> Self:
         """
         This method is overridden to use for BillingLineItem CPT codes.
         The codes are saved as string values in the BillingLineItem.cpt field,

canvas_sdk/v1/data/condition.py

@@ -1,7 +1,15 @@
+from typing import cast
+
 from django.db import models
 from django.db.models import TextChoices
 
-from canvas_sdk.v1.data.base import ValueSetLookupQuerySet
+from canvas_sdk.v1.data.base import (
+    BaseModelManager,
+    BaseQuerySet,
+    CommittableQuerySetMixin,
+    ForPatientQuerySetMixin,
+    ValueSetLookupQuerySetMixin,
+)
 
 
 class ClinicalStatus(TextChoices):
@@ -14,20 +22,28 @@ class ClinicalStatus(TextChoices):
     INVESTIGATIVE = "investigative", "investigative"
 
 
-class ConditionQuerySet(ValueSetLookupQuerySet):
+class ConditionQuerySet(
+    BaseQuerySet,
+    ValueSetLookupQuerySetMixin,
+    CommittableQuerySetMixin,
+    ForPatientQuerySetMixin,
+):
     """ConditionQuerySet."""
 
     pass
 
 
+ConditionManager = BaseModelManager.from_queryset(ConditionQuerySet)
+
+
 class Condition(models.Model):
     """Condition."""
 
     class Meta:
         managed = False
         db_table = "canvas_sdk_data_api_condition_001"
 
-    objects = ConditionQuerySet.as_manager()
+    objects = cast(ConditionQuerySet, ConditionManager())
 
     id = models.UUIDField()
     dbid = models.BigIntegerField(primary_key=True)

canvas_sdk/v1/data/lab.py

@@ -1,12 +1,26 @@
+from typing import cast
+
 from django.db import models
 
 from canvas_sdk.v1.data.base import (
-    CommittableModelManager,
+    BaseModelManager,
+    BaseQuerySet,
+    CommittableQuerySetMixin,
+    ForPatientQuerySetMixin,
     TimeframeLookupQuerySetMixin,
     ValueSetLookupQuerySet,
 )
 
 
+class LabReportQuerySet(BaseQuerySet, CommittableQuerySetMixin, ForPatientQuerySetMixin):
+    """A queryset for lab reports."""
+
+    pass
+
+
+LabReportManager = BaseModelManager.from_queryset(LabReportQuerySet)
+
+
 class TransmissionType(models.TextChoices):
     """Choices for transmission types."""
 
@@ -22,7 +36,7 @@ class Meta:
         managed = False
         db_table = "canvas_sdk_data_api_labreport_001"
 
-    objects = CommittableModelManager()
+    objects = cast(LabReportQuerySet, LabReportManager())
 
     id = models.UUIDField()
     dbid = models.BigIntegerField(primary_key=True)
@@ -52,14 +66,23 @@ class Meta:
     deleted = models.BooleanField()
 
 
+class LabReviewQuerySet(BaseQuerySet, CommittableQuerySetMixin, ForPatientQuerySetMixin):
+    """A queryset for lab reviews."""
+
+    pass
+
+
+LabReviewManager = BaseModelManager.from_queryset(LabReviewQuerySet)
+
+
 class LabReview(models.Model):
     """A class representing a lab review."""
 
     class Meta:
         managed = False
         db_table = "canvas_sdk_data_api_labreview_001"
 
-    objects = CommittableModelManager()
+    objects = cast(LabReviewQuerySet, LabReviewManager())
 
     id = models.UUIDField()
     dbid = models.BigIntegerField(primary_key=True)

canvas_sdk/v1/data/medication.py

@@ -1,7 +1,14 @@
+from typing import cast
+
 from django.db import models
 from django.db.models import TextChoices
 
-from canvas_sdk.v1.data.base import CommittableModelManager, ValueSetLookupQuerySet
+from canvas_sdk.v1.data.base import (
+    BaseModelManager,
+    CommittableQuerySetMixin,
+    ForPatientQuerySetMixin,
+    ValueSetLookupQuerySet,
+)
 
 
 class Status(TextChoices):
@@ -11,20 +18,23 @@ class Status(TextChoices):
     INACTIVE = "inactive", "inactive"
 
 
-class MedicationQuerySet(ValueSetLookupQuerySet):
+class MedicationQuerySet(ValueSetLookupQuerySet, CommittableQuerySetMixin, ForPatientQuerySetMixin):
     """MedicationQuerySet."""
 
     pass
 
 
+MedicationManager = BaseModelManager.from_queryset(MedicationQuerySet)
+
+
 class Medication(models.Model):
     """Medication."""
 
     class Meta:
         managed = False
         db_table = "canvas_sdk_data_api_medication_001"
 
-    objects = CommittableModelManager.from_queryset(MedicationQuerySet)()
+    objects = cast(MedicationQuerySet, MedicationManager())
 
     id = models.UUIDField()
     dbid = models.BigIntegerField(primary_key=True)

canvas_sdk/v1/data/observation.py

@@ -1,22 +1,34 @@
+from typing import cast
+
 from django.db import models
 
-from canvas_sdk.v1.data.base import CommittableModelManager, ValueSetLookupQuerySet
+from canvas_sdk.v1.data.base import (
+    BaseModelManager,
+    CommittableQuerySetMixin,
+    ForPatientQuerySetMixin,
+    ValueSetLookupQuerySet,
+)
 
 
-class ObservationQuerySet(ValueSetLookupQuerySet):
+class ObservationQuerySet(
+    ValueSetLookupQuerySet, CommittableQuerySetMixin, ForPatientQuerySetMixin
+):
     """ObservationQuerySet."""
 
     pass
 
 
+ObservationManager = BaseModelManager.from_queryset(ObservationQuerySet)
+
+
 class Observation(models.Model):
     """Observation."""
 
     class Meta:
         managed = False
         db_table = "canvas_sdk_data_api_observation_001"
 
-    objects = CommittableModelManager.from_queryset(ObservationQuerySet)()
+    objects = cast(ObservationQuerySet, ObservationManager())
 
     id = models.UUIDField()
     dbid = models.BigIntegerField(primary_key=True)

canvas_sdk/v1/data/protocol_override.py

@@ -1,6 +1,13 @@
+from typing import cast
+
 from django.db import models
 
-from canvas_sdk.v1.data.base import CommittableModelManager
+from canvas_sdk.v1.data.base import (
+    BaseModelManager,
+    BaseQuerySet,
+    CommittableQuerySetMixin,
+    ForPatientQuerySetMixin,
+)
 
 
 class IntervalUnit(models.TextChoices):
@@ -18,14 +25,23 @@ class Status(models.TextChoices):
     INACTIVE = "inactive", "inactive"
 
 
+class ProtocolOverrideQuerySet(BaseQuerySet, ForPatientQuerySetMixin, CommittableQuerySetMixin):
+    """ProtocolOverrideQuerySet."""
+
+    pass
+
+
+ProtocolOverrideManager = BaseModelManager.from_queryset(ProtocolOverrideQuerySet)
+
+
 class ProtocolOverride(models.Model):
     """ProtocolOverride."""
 
     class Meta:
         managed = False
         db_table = "canvas_sdk_data_api_protocoloverride_001"
 
-    objects = CommittableModelManager()
+    objects = cast(ProtocolOverrideQuerySet, ProtocolOverrideManager())
 
     id = models.UUIDField()
     dbid = models.BigIntegerField(primary_key=True)