prompt: 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()
  
  ```
  
golden_diff: 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()
  
verification_info: {"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"}]}