feat(firestore): add new pipeline expressions

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

@@ -48,6 +48,75 @@
 )
 
 
+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."""
 
@@ -90,6 +159,31 @@ def _to_pb(self) -> Value:
         )
 
 
+class Type(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.
@@ -120,6 +214,8 @@ def _cast_to_expr_or_convert_to_constant(
         """Convert arbitrary object to an Expression."""
         if isinstance(o, Expression):
             return o
+        if isinstance(o, Enum):
+            o = o.value
         if isinstance(o, dict):
             return Map(o)
         if isinstance(o, list):
@@ -2079,20 +2175,21 @@ def unix_seconds_to_timestamp(self) -> "Expression":
 
     @expose_as_static
     def timestamp_add(
-        self, unit: Expression | str, amount: Expression | float
+        self, unit: TimeUnit | str | Expression, amount: Expression | float
     ) -> "Expression":
         """Creates an expression that adds a specified amount of time to this timestamp expression.
 
         Example:
+            >>> # Add 1.5 days to the 'timestamp' field using TimeUnit enum.
+            >>> Field.of("timestamp").timestamp_add(TimeUnit.DAY, 1.5)
             >>> # Add a duration specified by the 'unit' and 'amount' fields to the 'timestamp' field.
             >>> Field.of("timestamp").timestamp_add(Field.of("unit"), Field.of("amount"))
-            >>> # Add 1.5 days to the 'timestamp' field.
+            >>> # Add 1.5 days to the 'timestamp' field using a string.
             >>> Field.of("timestamp").timestamp_add("day", 1.5)
 
         Args:
-            unit: The expression or string evaluating to the unit of time to add, must be one of
-                  'microsecond', 'millisecond', 'second', 'minute', 'hour', 'day'.
-            amount: The expression or float representing the amount of time to add.
+            unit: The unit of time to add.
+            amount: The amount of time to add.
 
         Returns:
             A new `Expression` representing the resulting timestamp.
@@ -2108,20 +2205,21 @@ def timestamp_add(
 
     @expose_as_static
     def timestamp_subtract(
-        self, unit: Expression | str, amount: Expression | float
+        self, unit: TimeUnit | str | Expression, amount: Expression | float
     ) -> "Expression":
         """Creates an expression that subtracts a specified amount of time from this timestamp expression.
 
         Example:
+            >>> # Subtract 2.5 hours from the 'timestamp' field using TimeUnit enum.
+            >>> Field.of("timestamp").timestamp_subtract(TimeUnit.HOUR, 2.5)
             >>> # Subtract a duration specified by the 'unit' and 'amount' fields from the 'timestamp' field.
             >>> Field.of("timestamp").timestamp_subtract(Field.of("unit"), Field.of("amount"))
-            >>> # Subtract 2.5 hours from the 'timestamp' field.
+            >>> # Subtract 2.5 hours from the 'timestamp' field using a string.
             >>> Field.of("timestamp").timestamp_subtract("hour", 2.5)
 
         Args:
-            unit: The expression or string evaluating to the unit of time to subtract, must be one of
-                  'microsecond', 'millisecond', 'second', 'minute', 'hour', 'day'.
-            amount: The expression or float representing the amount of time to subtract.
+            unit: The unit of time to subtract.
+            amount: The amount of time to subtract.
 
         Returns:
             A new `Expression` representing the resulting timestamp.
@@ -2206,6 +2304,163 @@ def as_(self, alias: str) -> "AliasedExpression":
         """
         return AliasedExpression(self, alias)
 
+    @expose_as_static
+    def cmp(self, other: Expression | CONSTANT_TYPE) -> "Expression":
+        """Creates an expression that compares this expression to another expression.
+
+        Returns an integer:
+            * -1 if this expression is less than the other
+            * 0 if they are equal
+            * 1 if this expression is greater than the other
+
+        Example:
+            >>> # Compare the 'price' field to 10
+            >>> Field.of("price").cmp(10)
+
+        Args:
+            other: The value to compare against.
+
+        Returns:
+            A new `Expression` representing the comparison operation.
+        """
+        return FunctionExpression(
+            "cmp", [self, self._cast_to_expr_or_convert_to_constant(other)]
+        )
+
+    @expose_as_static
+    def timestamp_trunc(
+        self,
+        granularity: TimeGranularity | Expression | str,
+        timezone: Expression | str | None = None,
+    ) -> "Expression":
+        """Creates an expression that truncates a timestamp to a specified granularity.
+
+        Example:
+            >>> # Truncate the 'createdAt' field to the day using TimeGranularity enum
+            >>> Field.of("createdAt").timestamp_trunc(TimeGranularity.DAY)
+            >>> # Truncate the 'createdAt' field to the day in the 'America/Los_Angeles' timezone
+            >>> Field.of("createdAt").timestamp_trunc(TimeGranularity.DAY, "America/Los_Angeles")
+            >>> # Truncate the 'createdAt' field to the day using a string
+            >>> Field.of("createdAt").timestamp_trunc("day")
+
+        Args:
+            granularity: The granularity to truncate to.
+            timezone: The optional timezone.
+
+        Returns:
+            A new `Expression` representing the timestamp_trunc operation.
+        """
+        args = [self, self._cast_to_expr_or_convert_to_constant(granularity)]
+        if timezone is not None:
+            args.append(self._cast_to_expr_or_convert_to_constant(timezone))
+        return FunctionExpression("timestamp_trunc", args)
+
+    @expose_as_static
+    def timestamp_extract(
+        self,
+        part: TimePart | str | Expression,
+        timezone: str | Expression | None = None,
+    ) -> "Expression":
+        """Creates an expression that extracts a part of a timestamp.
+
+        Example:
+            >>> # Extract the year from the 'createdAt' field using TimePart enum
+            >>> Field.of("createdAt").timestamp_extract(TimePart.YEAR)
+            >>> # Extract the year from the 'createdAt' field in the 'America/Los_Angeles' timezone
+            >>> Field.of("createdAt").timestamp_extract(TimePart.YEAR, "America/Los_Angeles")
+            >>> # Extract the year from the 'createdAt' field using a string
+            >>> Field.of("createdAt").timestamp_extract("year")
+
+        Args:
+            part: The part to extract.
+            timezone: The optional timezone.
+
+        Returns:
+            A new `Expression` representing the timestamp_extract operation.
+        """
+        args = [self, self._cast_to_expr_or_convert_to_constant(part)]
+        if timezone is not None:
+            args.append(self._cast_to_expr_or_convert_to_constant(timezone))
+        return FunctionExpression("timestamp_extract", args)
+
+    @expose_as_static
+    def timestamp_diff(
+        self, start: Expression | datetime.datetime, unit: TimeUnit | str | Expression
+    ) -> "Expression":
+        """Creates an expression that computes the difference between two timestamps in the specified unit.
+
+        Example:
+            >>> # Compute the difference in days between the 'end' field and the 'start' field using TimeUnit enum
+            >>> Field.of("end").timestamp_diff(Field.of("start"), TimeUnit.DAY)
+            >>> # Compute the difference in days using a string
+            >>> Field.of("end").timestamp_diff(Field.of("start"), "day")
+
+        Args:
+            start: The start timestamp.
+            unit: The unit of time.
+
+        Returns:
+            A new `Expression` representing the timestamp_diff operation.
+        """
+        return FunctionExpression(
+            "timestamp_diff",
+            [
+                self,
+                self._cast_to_expr_or_convert_to_constant(start),
+                self._cast_to_expr_or_convert_to_constant(unit),
+            ],
+        )
+
+    @expose_as_static
+    def if_null(self, *others: Expression | CONSTANT_TYPE) -> "Expression":
+        """Creates an expression that returns the first non-null expression from the provided arguments.
+
+        Example:
+            >>> # Return the 'nickname' field if not null, otherwise return 'firstName'
+            >>> Field.of("nickname").if_null(Field.of("firstName"))
+
+        Args:
+            *others: Additional expressions or constants to evaluate if the previous ones are null.
+
+        Returns:
+            A new `Expression` representing the if_null operation.
+        """
+        return FunctionExpression(
+            "if_null",
+            [self] + [self._cast_to_expr_or_convert_to_constant(o) for o in others],
+        )
+
+    @expose_as_static
+    def type(self) -> "Expression":
+        """Creates an expression that returns the data type of this expression's result as a string.
+
+        Example:
+            >>> # Get the type of the 'title' field
+            >>> Field.of("title").type()
+
+        Returns:
+            A new `Expression` representing the type operation.
+        """
+        return FunctionExpression("type", [self])
+
+    @expose_as_static
+    def is_type(self, type_val: Type | str | Expression) -> "BooleanExpression":
+        """Creates an expression that checks if the result is of the specified type.
+
+        Example:
+            >>> # Check if the 'price' field is a number
+            >>> Field.of("price").is_type("number")
+
+        Args:
+            type_val: The type string or expression to check against.
+
+        Returns:
+            A new `BooleanExpression` representing the is_type operation.
+        """
+        return BooleanExpression(
+            "is_type", [self, self._cast_to_expr_or_convert_to_constant(type_val)]
+        )
+
 
 class Constant(Expression, Generic[CONSTANT_TYPE]):
     """Represents a constant literal value in an expression."""