Make reference location extensible too

Author: greenape

flowmachine/flowmachine/core/server/query_schemas/active_at_reference_location_counts.py

@@ -15,13 +15,13 @@
     RedactedActiveAtReferenceLocationCounts,
 )
 
-from . import BaseExposedQuery
-
 __all__ = [
     "ActiveAtReferenceLocationCountsSchema",
     "ActiveAtReferenceLocationCountsExposed",
 ]
 
+from .base_exposed_query import BaseExposedQuery
+
 from .base_schema import BaseSchema
 
 from .reference_location import ReferenceLocationSchema

flowmachine/flowmachine/core/server/query_schemas/consecutive_trips_od_matrix.py

@@ -12,7 +12,7 @@
 from flowmachine.features.location.consecutive_trips_od_matrix import (
     ConsecutiveTripsODMatrix,
 )
-from . import BaseExposedQuery
+from .base_exposed_query import BaseExposedQuery
 from .base_schema import BaseSchema
 from .custom_fields import SubscriberSubset
 from .aggregation_unit import AggregationUnit, get_spatial_unit_obj

flowmachine/flowmachine/core/server/query_schemas/flowable_queries.py

@@ -4,6 +4,7 @@
 
 from flowmachine.core.server.query_schemas.daily_location import DailyLocationSchema
 from flowmachine.core.server.query_schemas.modal_location import ModalLocationSchema
+from flowmachine.core.server.query_schemas.unique_locations import UniqueLocationsSchema
 
 flowable_queries = {
     "daily_location": DailyLocationSchema,

flowmachine/flowmachine/core/server/query_schemas/joined_spatial_aggregate.py

@@ -12,9 +12,6 @@
 from marshmallow_oneofschema import OneOfSchema
 
 
-from flowmachine.core.server.query_schemas.spatial_aggregate import (
-    InputToSpatialAggregate,
-)
 from flowmachine.features.location.joined_spatial_aggregate import (
     JoinedSpatialAggregate,
 )
@@ -23,6 +20,7 @@
 )
 from .base_exposed_query import BaseExposedQuery
 from .base_schema import BaseSchema
+from .reference_location import ReferenceLocationSchema
 from .util import get_type_schemas_from_entrypoint
 
 
@@ -63,7 +61,7 @@ def _flowmachine_query_obj(self):
 class JoinedSpatialAggregateSchema(BaseSchema):
     # query_kind parameter is required here for claims validation
     query_kind = fields.String(validate=OneOf(["joined_spatial_aggregate"]))
-    locations = fields.Nested(InputToSpatialAggregate, required=True)
+    locations = fields.Nested(ReferenceLocationSchema, required=True)
     metric = fields.Nested(JoinableMetrics, required=True)
     method = fields.String(validate=OneOf(JoinedSpatialAggregate.allowed_methods))
 

flowmachine/flowmachine/core/server/query_schemas/reference_location.py

@@ -3,14 +3,9 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from marshmallow_oneofschema import OneOfSchema
-
-from flowmachine.core.server.query_schemas.daily_location import DailyLocationSchema
-from flowmachine.core.server.query_schemas.modal_location import ModalLocationSchema
+from flowmachine.core.server.query_schemas.util import get_type_schemas_from_entrypoint
 
 
 class ReferenceLocationSchema(OneOfSchema):
     type_field = "query_kind"
-    type_schemas = {
-        "daily_location": DailyLocationSchema,
-        "modal_location": ModalLocationSchema,
-    }
+    type_schemas = get_type_schemas_from_entrypoint("reference_location_queries")

flowmachine/flowmachine/core/server/query_schemas/reference_location_queries.py

@@ -0,0 +1,11 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from flowmachine.core.server.query_schemas.daily_location import DailyLocationSchema
+from flowmachine.core.server.query_schemas.modal_location import ModalLocationSchema
+
+reference_location_queries = {
+    "daily_location": DailyLocationSchema,
+    "modal_location": ModalLocationSchema,
+}

flowmachine/flowmachine/core/server/query_schemas/spatial_aggregate.py

@@ -2,34 +2,23 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marshmallow import fields, post_load
+from marshmallow import fields
 from marshmallow.validate import OneOf
-from marshmallow_oneofschema import OneOfSchema
 
 from flowmachine.features.location.redacted_spatial_aggregate import (
     RedactedSpatialAggregate,
 )
 from flowmachine.features.location.spatial_aggregate import SpatialAggregate
 from .base_exposed_query import BaseExposedQuery
 from .base_schema import BaseSchema
-from .daily_location import DailyLocationSchema
-from .modal_location import ModalLocationSchema
+from .reference_location import ReferenceLocationSchema
 
 __all__ = [
     "SpatialAggregateSchema",
     "SpatialAggregateExposed",
-    "InputToSpatialAggregate",
 ]
 
 
-class InputToSpatialAggregate(OneOfSchema):
-    type_field = "query_kind"
-    type_schemas = {
-        "daily_location": DailyLocationSchema,
-        "modal_location": ModalLocationSchema,
-    }
-
-
 class SpatialAggregateExposed(BaseExposedQuery):
     def __init__(self, *, locations):
         # Note: all input parameters need to be defined as attributes on `self`
@@ -54,6 +43,6 @@ def _flowmachine_query_obj(self):
 class SpatialAggregateSchema(BaseSchema):
     # query_kind parameter is required here for claims validation
     query_kind = fields.String(validate=OneOf(["spatial_aggregate"]))
-    locations = fields.Nested(InputToSpatialAggregate, required=True)
+    locations = fields.Nested(ReferenceLocationSchema, required=True)
 
     __model__ = SpatialAggregateExposed

flowmachine/flowmachine/core/server/query_schemas/trips_od_matrix.py

@@ -8,7 +8,7 @@
 from flowmachine.features.location.redacted_trips_od_matrix import RedactedTripsODMatrix
 from flowmachine.features.utilities.subscriber_locations import SubscriberLocations
 from flowmachine.features.location.trips_od_matrix import TripsODMatrix
-from . import BaseExposedQuery
+from .base_exposed_query import BaseExposedQuery
 from .base_schema import BaseSchema
 from .custom_fields import SubscriberSubset
 from .aggregation_unit import AggregationUnit, get_spatial_unit_obj

flowmachine/flowmachine/core/server/query_schemas/unique_visitor_counts.py

@@ -2,7 +2,7 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-from marshmallow import fields, post_load
+from marshmallow import fields
 from marshmallow.validate import OneOf
 
 from flowmachine.features.location.unique_visitor_counts import UniqueVisitorCounts
@@ -11,7 +11,7 @@
 
 from .unique_subscriber_counts import UniqueSubscriberCountsSchema
 
-from . import BaseExposedQuery
+from .base_exposed_query import BaseExposedQuery
 
 __all__ = [
     "UniqueVisitorCountsSchema",

flowmachine/flowmachine/core/server/query_schemas/unmoving_at_reference_location_counts.py

@@ -14,7 +14,7 @@
 from flowmachine.features.subscriber.unmoving_at_reference_location import (
     UnmovingAtReferenceLocation,
 )
-from . import BaseExposedQuery
+from .base_exposed_query import BaseExposedQuery
 from .base_schema import BaseSchema
 from .reference_location import ReferenceLocationSchema
 from .unique_locations import UniqueLocationsSchema