camel-agents / apps /common /auto_zip.py
DmitriiKhizbullin's picture
Proper sync
5ca24d3
raw
history blame
1.83 kB
# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
# Licensed under the Apache License, Version 2.0 (the “License”);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an “AS IS” BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
import json
import os
import zipfile
class AutoZip:
def __init__(self, zip_path: str, ext: str = ".json"):
self.zip_path = zip_path
self.zip = zipfile.ZipFile(zip_path, "r")
self.fl = [f for f in self.zip.filelist if f.filename.endswith(ext)]
def __next__(self):
if self.index >= len(self.fl):
raise StopIteration
else:
finfo = self.fl[self.index]
with self.zip.open(finfo) as f:
raw_json = json.loads(f.read().decode("utf-8"))
self.index += 1
return raw_json
def __len__(self):
return len(self.fl)
def __iter__(self):
self.index = 0
return self
def as_dict(self, include_zip_name: bool = False):
d = dict()
for finfo in self.fl:
with self.zip.open(finfo) as f:
raw_text = f.read().decode("utf-8")
if include_zip_name:
key = os.path.split(self.zip_path)[1] + "/" + finfo.filename
else:
key = finfo.filename
d[key] = raw_text
return d