diff --git a/gcsfs/core.py b/gcsfs/core.py index 8fe06625..e1e2bec1 100644 --- a/gcsfs/core.py +++ b/gcsfs/core.py @@ -1963,12 +1963,16 @@ def _upload_chunk(self, final=False): continue else: self.checker.update(data) + if final and contents: + j = json.loads(contents) + self.generation = j.get("generation") else: assert final, "Response looks like upload is over" if l: j = json.loads(contents) self.checker.update(data) self.checker.validate_json_response(j) + self.generation = j.get("generation") # Clear buffer and update offset when all is received self.buffer = UnclosableBytesIO() self.offset += l @@ -2014,7 +2018,7 @@ def _simple_upload(self): """One-shot upload, less than 5MB""" self.buffer.seek(0) data = self.buffer.read() - asyn.sync( + j = asyn.sync( self.gcsfs.loop, simple_upload, self.gcsfs, @@ -2029,6 +2033,7 @@ def _simple_upload(self): kms_key_name=self.kms_key_name, timeout=self.timeout, ) + self.generation = j.get("generation") def _fetch_range(self, start=None, end=None): """Get data from GCS @@ -2163,3 +2168,4 @@ async def simple_upload( ) checker.update(datain) checker.validate_json_response(j) + return j diff --git a/gcsfs/tests/test_core_versioned.py b/gcsfs/tests/test_core_versioned.py index 24821114..0b821acb 100644 --- a/gcsfs/tests/test_core_versioned.py +++ b/gcsfs/tests/test_core_versioned.py @@ -117,3 +117,19 @@ def test_find_versioned(gcs_versioned): versions = {f"{a}#{v1}", f"{a}#{v2}"} assert versions == set(gcs_versioned.find(a, versions=True)) assert versions == set(gcs_versioned.find(a, detail=True, versions=True)) + + +def test_write_captures_generation(gcs_versioned): + with gcs_versioned.open(a, "wb") as wo: + wo.write(b"test content") + assert wo.generation is not None + assert wo.generation == gcs_versioned.info(a)["generation"] + + +def test_write_captures_generation_multipart(gcs_versioned): + with gcs_versioned.open(b, "wb") as wo: + wo.write(b"first chunk") + wo.flush() + wo.write(b"second chunk") + assert wo.generation is not None + assert wo.generation == gcs_versioned.info(b)["generation"]