moved options and enums into pipeline_types

Author: daniel-sanche

packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py

@@ -18,6 +18,7 @@
 
 
 from google.cloud.firestore_v1 import pipeline_stages as stages
+from google.cloud.firestore_v1 import pipeline_types as types
 from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
 from google.cloud.firestore_v1.pipeline_expressions import (
     AggregateFunction,
@@ -275,7 +276,7 @@ def find_nearest(
         field: str | Expression,
         vector: Sequence[float] | "Vector",
         distance_measure: "DistanceMeasure",
-        options: stages.FindNearestOptions | None = None,
+        options: types.FindNearestOptions | None = None,
     ) -> "_BasePipeline":
         """
         Performs vector distance (similarity) search with given parameters on the
@@ -396,7 +397,7 @@ def sort(self, *orders: stages.Ordering) -> "_BasePipeline":
         return self._append(stages.Sort(*orders))
 
     def search(
-        self, query_or_options: str | BooleanExpression | stages.SearchOptions
+        self, query_or_options: str | BooleanExpression | types.SearchOptions
     ) -> "_BasePipeline":
         """
         Adds a search stage to the pipeline.
@@ -426,7 +427,7 @@ def search(
         """
         return self._append(stages.Search(query_or_options))
 
-    def sample(self, limit_or_options: int | stages.SampleOptions) -> "_BasePipeline":
+    def sample(self, limit_or_options: int | types.SampleOptions) -> "_BasePipeline":
         """
         Performs a pseudo-random sampling of the documents from the previous stage.
 
@@ -490,7 +491,7 @@ def unnest(
         self,
         field: str | Selectable,
         alias: str | Field | None = None,
-        options: stages.UnnestOptions | None = None,
+        options: types.UnnestOptions | None = None,
     ) -> "_BasePipeline":
         """
         Produces a document for each element in an array field from the previous stage document.

packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py

@@ -29,8 +29,15 @@
     from google.cloud.firestore_v1.base_pipeline import _BasePipeline
 
 from google.cloud.firestore_v1._helpers import GeoPoint, decode_value, encode_value
-from google.cloud.firestore_v1.types.document import Value
+from google.cloud.firestore_v1.pipeline_types import (
+    Ordering,
+    PipelineDataType,
+    TimeGranularity,
+    TimePart,
+    TimeUnit,
+)
 from google.cloud.firestore_v1.types.document import Pipeline as Pipeline_pb
+from google.cloud.firestore_v1.types.document import Value
 from google.cloud.firestore_v1.types.query import StructuredQuery as Query_pb
 from google.cloud.firestore_v1.vector import Vector
 
@@ -48,142 +55,6 @@
 )
 
 
-class TimeUnit(str, Enum):
-    """Enumeration of the different time units supported by the Firestore backend."""
-
-    MICROSECOND = "microsecond"
-    MILLISECOND = "millisecond"
-    SECOND = "second"
-    MINUTE = "minute"
-    HOUR = "hour"
-    DAY = "day"
-
-
-class TimeGranularity(str, Enum):
-    """Enumeration of the different time granularities supported by the Firestore backend."""
-
-    # Inherit from TimeUnit
-    MICROSECOND = TimeUnit.MICROSECOND.value
-    MILLISECOND = TimeUnit.MILLISECOND.value
-    SECOND = TimeUnit.SECOND.value
-    MINUTE = TimeUnit.MINUTE.value
-    HOUR = TimeUnit.HOUR.value
-    DAY = TimeUnit.DAY.value
-
-    # Additional granularities
-    WEEK = "week"
-    WEEK_MONDAY = "week(monday)"
-    WEEK_TUESDAY = "week(tuesday)"
-    WEEK_WEDNESDAY = "week(wednesday)"
-    WEEK_THURSDAY = "week(thursday)"
-    WEEK_FRIDAY = "week(friday)"
-    WEEK_SATURDAY = "week(saturday)"
-    WEEK_SUNDAY = "week(sunday)"
-    ISOWEEK = "isoweek"
-    MONTH = "month"
-    QUARTER = "quarter"
-    YEAR = "year"
-    ISOYEAR = "isoyear"
-
-
-class TimePart(str, Enum):
-    """Enumeration of the different time parts supported by the Firestore backend."""
-
-    # Inherit from TimeUnit
-    MICROSECOND = TimeUnit.MICROSECOND.value
-    MILLISECOND = TimeUnit.MILLISECOND.value
-    SECOND = TimeUnit.SECOND.value
-    MINUTE = TimeUnit.MINUTE.value
-    HOUR = TimeUnit.HOUR.value
-    DAY = TimeUnit.DAY.value
-
-    # Inherit from TimeGranularity
-    WEEK = TimeGranularity.WEEK.value
-    WEEK_MONDAY = TimeGranularity.WEEK_MONDAY.value
-    WEEK_TUESDAY = TimeGranularity.WEEK_TUESDAY.value
-    WEEK_WEDNESDAY = TimeGranularity.WEEK_WEDNESDAY.value
-    WEEK_THURSDAY = TimeGranularity.WEEK_THURSDAY.value
-    WEEK_FRIDAY = TimeGranularity.WEEK_FRIDAY.value
-    WEEK_SATURDAY = TimeGranularity.WEEK_SATURDAY.value
-    WEEK_SUNDAY = TimeGranularity.WEEK_SUNDAY.value
-    ISOWEEK = TimeGranularity.ISOWEEK.value
-    MONTH = TimeGranularity.MONTH.value
-    QUARTER = TimeGranularity.QUARTER.value
-    YEAR = TimeGranularity.YEAR.value
-    ISOYEAR = TimeGranularity.ISOYEAR.value
-
-    # Additional parts
-    DAY_OF_WEEK = "dayofweek"
-    DAY_OF_YEAR = "dayofyear"
-
-
-class Ordering:
-    """Represents the direction for sorting results in a pipeline."""
-
-    class Direction(Enum):
-        ASCENDING = "ascending"
-        DESCENDING = "descending"
-
-    def __init__(self, expr, order_dir: Direction | str = Direction.ASCENDING):
-        """
-        Initializes an Ordering instance
-
-        Args:
-            expr (Expression | str): The expression or field path string to sort by.
-                If a string is provided, it's treated as a field path.
-            order_dir (Direction | str): The direction to sort in.
-                Defaults to ascending
-        """
-        self.expr = expr if isinstance(expr, Expression) else Field.of(expr)
-        self.order_dir = (
-            Ordering.Direction[order_dir.upper()]
-            if isinstance(order_dir, str)
-            else order_dir
-        )
-
-    def __repr__(self):
-        if self.order_dir is Ordering.Direction.ASCENDING:
-            order_str = ".ascending()"
-        else:
-            order_str = ".descending()"
-        return f"{self.expr!r}{order_str}"
-
-    def _to_pb(self) -> Value:
-        return Value(
-            map_value={
-                "fields": {
-                    "direction": Value(string_value=self.order_dir.value),
-                    "expression": self.expr._to_pb(),
-                }
-            }
-        )
-
-
-class PipelineDataType(str, Enum):
-    """Enumeration of the different types generated by the Firestore backend."""
-
-    NULL = "null"
-    ARRAY = "array"
-    BOOLEAN = "boolean"
-    BYTES = "bytes"
-    TIMESTAMP = "timestamp"
-    GEO_POINT = "geo_point"
-    NUMBER = "number"
-    INT32 = "int32"
-    INT64 = "int64"
-    FLOAT64 = "float64"
-    DECIMAL128 = "decimal128"
-    MAP = "map"
-    REFERENCE = "reference"
-    STRING = "string"
-    VECTOR = "vector"
-    MAX_KEY = "max_key"
-    MIN_KEY = "min_key"
-    OBJECT_ID = "object_id"
-    REGEX = "regex"
-    REQUEST_TIMESTAMP = "request_timestamp"
-
-
 class Expression(ABC):
     """Represents an expression that can be evaluated to a value within the
     execution of a pipeline.
@@ -2503,7 +2374,9 @@ def type(self) -> "Expression":
         return FunctionExpression("type", [self])
 
     @expose_as_static
-    def is_type(self, type_val: PipelineDataType | str | Expression) -> "BooleanExpression":
+    def is_type(
+        self, type_val: PipelineDataType | str | Expression
+    ) -> "BooleanExpression":
         """Creates an expression that checks if the result is of the specified type.
 
         Example: