We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Firestore] WriteBatch doesn't return instance so you cannot chain.
The WriteBatch methods don’t return the WriteBatch instances for chaining.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `firestore/google/cloud/firestore_v1beta1/batch.py`
Content:
```
1 # Copyright 2017 Google LLC All rights reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Helpers for batch requests to the Google Cloud Firestore API."""
16 
17 
18 from google.cloud.firestore_v1beta1 import _helpers
19 
20 
21 class WriteBatch(object):
22     """Accumulate write operations to be sent in a batch.
23 
24     This has the same set of methods for write operations that
25     :class:`~.firestore_v1beta1.document.DocumentReference` does,
26     e.g. :meth:`~.firestore_v1beta1.document.DocumentReference.create`.
27 
28     Args:
29         client (~.firestore_v1beta1.client.Client): The client that
30             created this batch.
31     """
32 
33     def __init__(self, client):
34         self._client = client
35         self._write_pbs = []
36 
37     def _add_write_pbs(self, write_pbs):
38         """Add `Write`` protobufs to this transaction.
39 
40         This method intended to be over-ridden by subclasses.
41 
42         Args:
43             write_pbs (List[google.cloud.proto.firestore.v1beta1.\
44                 write_pb2.Write]): A list of write protobufs to be added.
45         """
46         self._write_pbs.extend(write_pbs)
47 
48     def create(self, reference, document_data):
49         """Add a "change" to this batch to create a document.
50 
51         If the document given by ``reference`` already exists, then this
52         batch will fail when :meth:`commit`-ed.
53 
54         Args:
55             reference (~.firestore_v1beta1.document.DocumentReference): A
56                 document reference to be created in this batch.
57             document_data (dict): Property names and values to use for
58                 creating a document.
59         """
60         write_pbs = _helpers.pbs_for_create(reference._document_path, document_data)
61         self._add_write_pbs(write_pbs)
62 
63     def set(self, reference, document_data, merge=False):
64         """Add a "change" to replace a document.
65 
66         See
67         :meth:`~.firestore_v1beta1.document.DocumentReference.set` for
68         more information on how ``option`` determines how the change is
69         applied.
70 
71         Args:
72             reference (~.firestore_v1beta1.document.DocumentReference):
73                 A document reference that will have values set in this batch.
74             document_data (dict):
75                 Property names and values to use for replacing a document.
76             merge (Optional[bool] or Optional[List<apispec>]):
77                 If True, apply merging instead of overwriting the state
78                 of the document.
79         """
80         if merge is not False:
81             write_pbs = _helpers.pbs_for_set_with_merge(
82                 reference._document_path, document_data, merge
83             )
84         else:
85             write_pbs = _helpers.pbs_for_set_no_merge(
86                 reference._document_path, document_data
87             )
88 
89         self._add_write_pbs(write_pbs)
90 
91     def update(self, reference, field_updates, option=None):
92         """Add a "change" to update a document.
93 
94         See
95         :meth:`~.firestore_v1beta1.document.DocumentReference.update` for
96         more information on ``field_updates`` and ``option``.
97 
98         Args:
99             reference (~.firestore_v1beta1.document.DocumentReference): A
100                document reference that will be deleted in this batch.
101             field_updates (dict): Field names or paths to update and values
102                 to update with.
103             option (Optional[~.firestore_v1beta1.client.WriteOption]): A
104                write option to make assertions / preconditions on the server
105                state of the document before applying changes.
106         """
107         if option.__class__.__name__ == "ExistsOption":
108             raise ValueError("you must not pass an explicit write option to " "update.")
109         write_pbs = _helpers.pbs_for_update(
110             reference._document_path, field_updates, option
111         )
112         self._add_write_pbs(write_pbs)
113 
114     def delete(self, reference, option=None):
115         """Add a "change" to delete a document.
116 
117         See
118         :meth:`~.firestore_v1beta1.document.DocumentReference.delete` for
119         more information on how ``option`` determines how the change is
120         applied.
121 
122         Args:
123             reference (~.firestore_v1beta1.document.DocumentReference): A
124                document reference that will be deleted in this batch.
125             option (Optional[~.firestore_v1beta1.client.WriteOption]): A
126                write option to make assertions / preconditions on the server
127                state of the document before applying changes.
128         """
129         write_pb = _helpers.pb_for_delete(reference._document_path, option)
130         self._add_write_pbs([write_pb])
131 
132     def commit(self):
133         """Commit the changes accumulated in this batch.
134 
135         Returns:
136             List[google.cloud.proto.firestore.v1beta1.\
137                 write_pb2.WriteResult, ...]: The write results corresponding
138             to the changes committed, returned in the same order as the
139             changes were applied to this batch. A write result contains an
140             ``update_time`` field.
141         """
142         commit_response = self._client._firestore_api.commit(
143             self._client._database_string,
144             self._write_pbs,
145             transaction=None,
146             metadata=self._client._rpc_metadata,
147         )
148 
149         self._write_pbs = []
150         return list(commit_response.write_results)
151
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/firestore/google/cloud/firestore_v1beta1/batch.py b/firestore/google/cloud/firestore_v1beta1/batch.py
--- a/firestore/google/cloud/firestore_v1beta1/batch.py
+++ b/firestore/google/cloud/firestore_v1beta1/batch.py
@@ -33,6 +33,8 @@
     def __init__(self, client):
         self._client = client
         self._write_pbs = []
+        self.write_results = None
+        self.commit_time = None
 
     def _add_write_pbs(self, write_pbs):
         """Add `Write`` protobufs to this transaction.
@@ -147,4 +149,13 @@
         )
 
         self._write_pbs = []
-        return list(commit_response.write_results)
+        self.write_results = results = list(commit_response.write_results)
+        self.commit_time = commit_response.commit_time
+        return results
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            self.commit()

{"golden_diff": "diff --git a/firestore/google/cloud/firestore_v1beta1/batch.py b/firestore/google/cloud/firestore_v1beta1/batch.py\n--- a/firestore/google/cloud/firestore_v1beta1/batch.py\n+++ b/firestore/google/cloud/firestore_v1beta1/batch.py\n@@ -33,6 +33,8 @@\n     def __init__(self, client):\n         self._client = client\n         self._write_pbs = []\n+        self.write_results = None\n+        self.commit_time = None\n \n     def _add_write_pbs(self, write_pbs):\n         \"\"\"Add `Write`` protobufs to this transaction.\n@@ -147,4 +149,13 @@\n         )\n \n         self._write_pbs = []\n-        return list(commit_response.write_results)\n+        self.write_results = results = list(commit_response.write_results)\n+        self.commit_time = commit_response.commit_time\n+        return results\n+\n+    def __enter__(self):\n+        return self\n+\n+    def __exit__(self, exc_type, exc_value, traceback):\n+        if exc_type is None:\n+            self.commit()\n", "issue": "[Firestore] WriteBatch doesn't return instance so you cannot chain.\nThe WriteBatch methods don\u2019t return the WriteBatch instances for chaining.\r\n\n", "before_files": [{"content": "# Copyright 2017 Google LLC All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helpers for batch requests to the Google Cloud Firestore API.\"\"\"\n\n\nfrom google.cloud.firestore_v1beta1 import _helpers\n\n\nclass WriteBatch(object):\n    \"\"\"Accumulate write operations to be sent in a batch.\n\n    This has the same set of methods for write operations that\n    :class:`~.firestore_v1beta1.document.DocumentReference` does,\n    e.g. :meth:`~.firestore_v1beta1.document.DocumentReference.create`.\n\n    Args:\n        client (~.firestore_v1beta1.client.Client): The client that\n            created this batch.\n    \"\"\"\n\n    def __init__(self, client):\n        self._client = client\n        self._write_pbs = []\n\n    def _add_write_pbs(self, write_pbs):\n        \"\"\"Add `Write`` protobufs to this transaction.\n\n        This method intended to be over-ridden by subclasses.\n\n        Args:\n            write_pbs (List[google.cloud.proto.firestore.v1beta1.\\\n                write_pb2.Write]): A list of write protobufs to be added.\n        \"\"\"\n        self._write_pbs.extend(write_pbs)\n\n    def create(self, reference, document_data):\n        \"\"\"Add a \"change\" to this batch to create a document.\n\n        If the document given by ``reference`` already exists, then this\n        batch will fail when :meth:`commit`-ed.\n\n        Args:\n            reference (~.firestore_v1beta1.document.DocumentReference): A\n                document reference to be created in this batch.\n            document_data (dict): Property names and values to use for\n                creating a document.\n        \"\"\"\n        write_pbs = _helpers.pbs_for_create(reference._document_path, document_data)\n        self._add_write_pbs(write_pbs)\n\n    def set(self, reference, document_data, merge=False):\n        \"\"\"Add a \"change\" to replace a document.\n\n        See\n        :meth:`~.firestore_v1beta1.document.DocumentReference.set` for\n        more information on how ``option`` determines how the change is\n        applied.\n\n        Args:\n            reference (~.firestore_v1beta1.document.DocumentReference):\n                A document reference that will have values set in this batch.\n            document_data (dict):\n                Property names and values to use for replacing a document.\n            merge (Optional[bool] or Optional[List<apispec>]):\n                If True, apply merging instead of overwriting the state\n                of the document.\n        \"\"\"\n        if merge is not False:\n            write_pbs = _helpers.pbs_for_set_with_merge(\n                reference._document_path, document_data, merge\n            )\n        else:\n            write_pbs = _helpers.pbs_for_set_no_merge(\n                reference._document_path, document_data\n            )\n\n        self._add_write_pbs(write_pbs)\n\n    def update(self, reference, field_updates, option=None):\n        \"\"\"Add a \"change\" to update a document.\n\n        See\n        :meth:`~.firestore_v1beta1.document.DocumentReference.update` for\n        more information on ``field_updates`` and ``option``.\n\n        Args:\n            reference (~.firestore_v1beta1.document.DocumentReference): A\n               document reference that will be deleted in this batch.\n            field_updates (dict): Field names or paths to update and values\n                to update with.\n            option (Optional[~.firestore_v1beta1.client.WriteOption]): A\n               write option to make assertions / preconditions on the server\n               state of the document before applying changes.\n        \"\"\"\n        if option.__class__.__name__ == \"ExistsOption\":\n            raise ValueError(\"you must not pass an explicit write option to \" \"update.\")\n        write_pbs = _helpers.pbs_for_update(\n            reference._document_path, field_updates, option\n        )\n        self._add_write_pbs(write_pbs)\n\n    def delete(self, reference, option=None):\n        \"\"\"Add a \"change\" to delete a document.\n\n        See\n        :meth:`~.firestore_v1beta1.document.DocumentReference.delete` for\n        more information on how ``option`` determines how the change is\n        applied.\n\n        Args:\n            reference (~.firestore_v1beta1.document.DocumentReference): A\n               document reference that will be deleted in this batch.\n            option (Optional[~.firestore_v1beta1.client.WriteOption]): A\n               write option to make assertions / preconditions on the server\n               state of the document before applying changes.\n        \"\"\"\n        write_pb = _helpers.pb_for_delete(reference._document_path, option)\n        self._add_write_pbs([write_pb])\n\n    def commit(self):\n        \"\"\"Commit the changes accumulated in this batch.\n\n        Returns:\n            List[google.cloud.proto.firestore.v1beta1.\\\n                write_pb2.WriteResult, ...]: The write results corresponding\n            to the changes committed, returned in the same order as the\n            changes were applied to this batch. A write result contains an\n            ``update_time`` field.\n        \"\"\"\n        commit_response = self._client._firestore_api.commit(\n            self._client._database_string,\n            self._write_pbs,\n            transaction=None,\n            metadata=self._client._rpc_metadata,\n        )\n\n        self._write_pbs = []\n        return list(commit_response.write_results)\n", "path": "firestore/google/cloud/firestore_v1beta1/batch.py"}], "after_files": [{"content": "# Copyright 2017 Google LLC All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helpers for batch requests to the Google Cloud Firestore API.\"\"\"\n\n\nfrom google.cloud.firestore_v1beta1 import _helpers\n\n\nclass WriteBatch(object):\n    \"\"\"Accumulate write operations to be sent in a batch.\n\n    This has the same set of methods for write operations that\n    :class:`~.firestore_v1beta1.document.DocumentReference` does,\n    e.g. :meth:`~.firestore_v1beta1.document.DocumentReference.create`.\n\n    Args:\n        client (~.firestore_v1beta1.client.Client): The client that\n            created this batch.\n    \"\"\"\n\n    def __init__(self, client):\n        self._client = client\n        self._write_pbs = []\n        self.write_results = None\n        self.commit_time = None\n\n    def _add_write_pbs(self, write_pbs):\n        \"\"\"Add `Write`` protobufs to this transaction.\n\n        This method intended to be over-ridden by subclasses.\n\n        Args:\n            write_pbs (List[google.cloud.proto.firestore.v1beta1.\\\n                write_pb2.Write]): A list of write protobufs to be added.\n        \"\"\"\n        self._write_pbs.extend(write_pbs)\n\n    def create(self, reference, document_data):\n        \"\"\"Add a \"change\" to this batch to create a document.\n\n        If the document given by ``reference`` already exists, then this\n        batch will fail when :meth:`commit`-ed.\n\n        Args:\n            reference (~.firestore_v1beta1.document.DocumentReference): A\n                document reference to be created in this batch.\n            document_data (dict): Property names and values to use for\n                creating a document.\n        \"\"\"\n        write_pbs = _helpers.pbs_for_create(reference._document_path, document_data)\n        self._add_write_pbs(write_pbs)\n\n    def set(self, reference, document_data, merge=False):\n        \"\"\"Add a \"change\" to replace a document.\n\n        See\n        :meth:`~.firestore_v1beta1.document.DocumentReference.set` for\n        more information on how ``option`` determines how the change is\n        applied.\n\n        Args:\n            reference (~.firestore_v1beta1.document.DocumentReference):\n                A document reference that will have values set in this batch.\n            document_data (dict):\n                Property names and values to use for replacing a document.\n            merge (Optional[bool] or Optional[List<apispec>]):\n                If True, apply merging instead of overwriting the state\n                of the document.\n        \"\"\"\n        if merge is not False:\n            write_pbs = _helpers.pbs_for_set_with_merge(\n                reference._document_path, document_data, merge\n            )\n        else:\n            write_pbs = _helpers.pbs_for_set_no_merge(\n                reference._document_path, document_data\n            )\n\n        self._add_write_pbs(write_pbs)\n\n    def update(self, reference, field_updates, option=None):\n        \"\"\"Add a \"change\" to update a document.\n\n        See\n        :meth:`~.firestore_v1beta1.document.DocumentReference.update` for\n        more information on ``field_updates`` and ``option``.\n\n        Args:\n            reference (~.firestore_v1beta1.document.DocumentReference): A\n               document reference that will be deleted in this batch.\n            field_updates (dict): Field names or paths to update and values\n                to update with.\n            option (Optional[~.firestore_v1beta1.client.WriteOption]): A\n               write option to make assertions / preconditions on the server\n               state of the document before applying changes.\n        \"\"\"\n        if option.__class__.__name__ == \"ExistsOption\":\n            raise ValueError(\"you must not pass an explicit write option to \" \"update.\")\n        write_pbs = _helpers.pbs_for_update(\n            reference._document_path, field_updates, option\n        )\n        self._add_write_pbs(write_pbs)\n\n    def delete(self, reference, option=None):\n        \"\"\"Add a \"change\" to delete a document.\n\n        See\n        :meth:`~.firestore_v1beta1.document.DocumentReference.delete` for\n        more information on how ``option`` determines how the change is\n        applied.\n\n        Args:\n            reference (~.firestore_v1beta1.document.DocumentReference): A\n               document reference that will be deleted in this batch.\n            option (Optional[~.firestore_v1beta1.client.WriteOption]): A\n               write option to make assertions / preconditions on the server\n               state of the document before applying changes.\n        \"\"\"\n        write_pb = _helpers.pb_for_delete(reference._document_path, option)\n        self._add_write_pbs([write_pb])\n\n    def commit(self):\n        \"\"\"Commit the changes accumulated in this batch.\n\n        Returns:\n            List[google.cloud.proto.firestore.v1beta1.\\\n                write_pb2.WriteResult, ...]: The write results corresponding\n            to the changes committed, returned in the same order as the\n            changes were applied to this batch. A write result contains an\n            ``update_time`` field.\n        \"\"\"\n        commit_response = self._client._firestore_api.commit(\n            self._client._database_string,\n            self._write_pbs,\n            transaction=None,\n            metadata=self._client._rpc_metadata,\n        )\n\n        self._write_pbs = []\n        self.write_results = results = list(commit_response.write_results)\n        self.commit_time = commit_response.commit_time\n        return results\n\n    def __enter__(self):\n        return self\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        if exc_type is None:\n            self.commit()\n", "path": "firestore/google/cloud/firestore_v1beta1/batch.py"}]}