Skip to content

Commit

Permalink
Fix s3 decoding
Browse files Browse the repository at this point in the history
  • Loading branch information
munrojm committed Feb 28, 2024
1 parent 83f421c commit a14719e
Showing 1 changed file with 8 additions and 5 deletions.
13 changes: 8 additions & 5 deletions mp_api/client/core/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -394,12 +394,10 @@ def _query_open_data(
)

if "jsonl" in key:
data = str([doc.strip() for doc in file.read().strip().split("\n")])
data = data.replace("'", "").replace("\\", "\\\\")
decoded_data = [decoder(jline) for jline in file.read().splitlines()]
else:
data = file.read()
decoded_data = decoder(file.read())

decoded_data = decoder(data)
decoded_data = (
[decoded_data] if not isinstance(decoded_data, list) else decoded_data
)
Expand Down Expand Up @@ -466,6 +464,7 @@ def _query_resource(
)
and num_chunks is None
and "substrates" not in self.suffix
and "phonon" not in self.suffix
)

if fields:
Expand Down Expand Up @@ -504,7 +503,11 @@ def _query_resource(
f"{suffix}" if is_tasks else f"collections/{db_version}/{suffix}"
)
objects = self.s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)
keys = [doc["Key"] for doc in objects["Contents"]]
keys = (
[doc["Key"] for doc in objects["Contents"]]
if "Contents" in objects
else []
)

decoder = MontyDecoder().decode if self.monty_decode else json.loads

Expand Down

0 comments on commit a14719e

Please sign in to comment.