feat: literals pipeline stage

google/cloud/firestore_v1/base_pipeline.py

@@ -273,6 +273,71 @@ def find_nearest(
             stages.FindNearest(field, vector, distance_measure, options)
         )
 
+    def literals(self, *documents: str | Selectable) -> "_BasePipeline":
+        """
+        Returns documents from a fixed set of predefined document objects.
+
+        This stage is commonly used for testing other stages in isolation, though it can
+        also be used as inputs to join conditions.
+
+        Example:
+            >>> from google.cloud.firestore_v1.pipeline_expressions import Constant
+            >>> documents = [
+            ...     {"name": "joe", "age": 10},
+            ...     {"name": "bob", "age": 30},
+            ...     {"name": "alice", "age": 40}
+            ... ]
+            >>> pipeline = client.pipeline()
+            ...     .literals(Constant.of(documents))
+            ...     .where(field("age").lessThan(35))
+
+            Output documents:
+            ```json
+            [
+                {"name": "joe", "age": 10},
+                {"name": "bob", "age": 30}
+            ]
+            ```
+
+        Behavior:
+            The `literals(...)` stage can only be used as the first stage in a pipeline (or
+            sub-pipeline). The order of documents returned from the `literals` matches the
+            order in which they are defined.
+
+            While literal values are the most common, it is also possible to pass in
+            expressions, which will be evaluated and returned, making it possible to test
+            out different query / expression behavior without first needing to create some
+            test data.
+
+            For example, the following shows how to quickly test out the `length(...)`
+            function on some constant test sets:
+
+        Example:
+            >>> from google.cloud.firestore_v1.pipeline_expressions import Constant
+            >>> documents = [
+            ...     {"x": Constant.of("foo-bar-baz").char_length()},
+            ...     {"x": Constant.of("bar").char_length()}
+            ... ]
+            >>> pipeline = client.pipeline().literals(Constant.of(documents))
+
+            Output documents:
+            ```json
+            [
+                {"x": 11},
+                {"x": 3}
+            ]
+            ```
+
+        Args:
+            documents: A `str` or `Selectable` expression. If a `str`, it's
+                       treated as a field path to an array of documents.
+                       If a `Selectable`, it's usually a `Constant`
+                       containing an array of documents (as dictionaries).
+        Returns:
+            A new Pipeline object with this stage appended to the stage list.
+        """
+        return self._append(stages.Literals(*documents))
+
     def replace_with(
         self,
         field: Selectable,

google/cloud/firestore_v1/pipeline_stages.py

@@ -342,6 +342,23 @@ def _pb_args(self):
         return [Value(integer_value=self.limit)]
 
 
+class Literals(Stage):
+    """Returns documents from a fixed set of predefined document objects."""
+
+    def __init__(self, *documents: str | Selectable):
+        super().__init__("literals")
+        self.documents = documents
+
+    def _pb_args(self):
+        args = []
+        for doc in self.documents:
+            if hasattr(doc, "_to_pb"):
+                args.append(doc._to_pb())
+            else:
+                args.append(encode_value(doc))
+        return args
+
+
 class Offset(Stage):
     """Skips a specified number of documents."""
 

tests/system/pipeline_e2e/general.yaml

@@ -684,4 +684,24 @@ tests:
         - args:
           - fieldReferenceValue: awards
           - stringValue: full_replace
-          name: replace_with
+          name: replace_with
+  - description: literals
+    pipeline:
+      - Literals:
+        - title: "The Hitchhiker's Guide to the Galaxy"
+          author: "Douglas Adams"
+    assert_results:
+      - title: "The Hitchhiker's Guide to the Galaxy"
+        author: "Douglas Adams"
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - mapValue:
+              fields:
+                author:
+                  stringValue: "Douglas Adams"
+                title:
+                  stringValue: "The Hitchhiker's Guide to the Galaxy"
+          name: literals
+

tests/unit/v1/test_pipeline.py

@@ -392,6 +392,7 @@ def test_pipeline_execute_stream_equivalence():
         ),
         ("replace_with", ("name",), stages.ReplaceWith),
         ("replace_with", (Field.of("n"),), stages.ReplaceWith),
+        ("literals", (Field.of("a"),), stages.Literals),
         ("sort", (Field.of("n").descending(),), stages.Sort),
         ("sort", (Field.of("n").descending(), Field.of("m").ascending()), stages.Sort),
         ("sample", (10,), stages.Sample),

tests/unit/v1/test_pipeline_stages.py

@@ -516,6 +516,35 @@ def test_to_pb(self):
         assert len(result.options) == 0
 
 
+class TestLiterals:
+    def _make_one(self, *args, **kwargs):
+        return stages.Literals(*args, **kwargs)
+
+    def test_ctor(self):
+        val1 = Constant.of({"a": 1})
+        val2 = Constant.of({"b": 2})
+        instance = self._make_one(val1, val2)
+        assert instance.documents == (val1, val2)
+        assert instance.name == "literals"
+
+    def test_repr(self):
+        val1 = Constant.of({"a": 1})
+        instance = self._make_one(val1)
+        repr_str = repr(instance)
+        assert repr_str == "Literals(documents=(Constant.of({'a': 1}),))"
+
+    def test_to_pb(self):
+        val1 = Constant.of({"a": 1})
+        val2 = Constant.of({"b": 2})
+        instance = self._make_one(val1, val2)
+        result = instance._to_pb()
+        assert result.name == "literals"
+        assert len(result.args) == 2
+        assert result.args[0].map_value.fields["a"].integer_value == 1
+        assert result.args[1].map_value.fields["b"].integer_value == 2
+        assert len(result.options) == 0
+
+
 class TestOffset:
     def _make_one(self, *args, **kwargs):
         return stages.Offset(*args, **kwargs)