|
7 | 7 | import datetime |
8 | 8 | import json |
9 | 9 | import logging |
| 10 | +import os |
| 11 | +import random |
10 | 12 | import re |
11 | 13 | import traceback |
12 | 14 | import zipfile |
|
43 | 45 | from .data.datasets_io import ( |
44 | 46 | read_external_dataset_array, |
45 | 47 | ) |
| 48 | +from .exception import UnparsableFile |
46 | 49 | from .introspection import ( |
47 | 50 | get_class_from_content_type, |
48 | 51 | get_obj_type, |
|
69 | 72 | serialize_xml, |
70 | 73 | read_energyml_xml_str, |
71 | 74 | read_energyml_xml_bytes, |
72 | | - read_energyml_json_str, |
| 75 | + read_energyml_json_str, read_energyml_json_bytes, JSON_VERSION, |
73 | 76 | ) |
74 | 77 | from .workspace import EnergymlWorkspace |
75 | 78 | from .xml import is_energyml_content_type |
@@ -130,6 +133,71 @@ def __str__(self): |
130 | 133 | # + f"\n{[serialize_json(ar) for ar in self.additional_rels]}" |
131 | 134 | ) |
132 | 135 |
|
| 136 | + def add_file(self, obj: Union[List, bytes, BytesIO, str, RawFile]): |
| 137 | + """ |
| 138 | + Add one ore multiple files to the epc file. |
| 139 | + For non energyml file, it is better to use the RawFile class. |
| 140 | + The input can be a single file content, file path, or a list of them |
| 141 | + :param obj: |
| 142 | + :return: |
| 143 | + """ |
| 144 | + if isinstance(obj, list): |
| 145 | + for o in obj: |
| 146 | + self.add_file(o) |
| 147 | + elif isinstance(obj, bytes) or isinstance(obj, BytesIO): |
| 148 | + try: |
| 149 | + xml_obj = read_energyml_xml_bytes(obj) |
| 150 | + self.energyml_objects.append(xml_obj) |
| 151 | + except: |
| 152 | + try: |
| 153 | + if isinstance(obj, BytesIO): |
| 154 | + obj.seek(0) |
| 155 | + json_obj = read_energyml_json_bytes(obj, json_version=JSON_VERSION.OSDU_OFFICIAL) |
| 156 | + self.add_file(json_obj) |
| 157 | + except: |
| 158 | + # if isinstance(obj, BytesIO): |
| 159 | + # obj.seek(0) |
| 160 | + # self.add_file(RawFile(path=f"pleaseRenameThisFile_{str(random.random())}", content=obj)) |
| 161 | + raise UnparsableFile() |
| 162 | + elif isinstance(obj, RawFile): |
| 163 | + self.raw_files.append(obj) |
| 164 | + elif isinstance(obj, str): |
| 165 | + # Can be a path or a content |
| 166 | + if os.path.exists(obj): |
| 167 | + with open(obj, "rb") as f: |
| 168 | + file_content = f.read() |
| 169 | + f_name = os.path.basename(obj) |
| 170 | + _, f_ext = os.path.splitext(f_name) |
| 171 | + if f_ext.lower().endswith(".xml") or f_ext.lower().endswith(".json"): |
| 172 | + try: |
| 173 | + self.add_file(file_content) |
| 174 | + except UnparsableFile: |
| 175 | + self.add_file(RawFile(f_name, BytesIO(file_content))) |
| 176 | + elif not f_ext.lower().endswith(".rels"): |
| 177 | + self.add_file(RawFile(f_name, BytesIO(file_content))) |
| 178 | + else: |
| 179 | + logging.error(f"Not supported file extension {f_name}") |
| 180 | + else: |
| 181 | + try: |
| 182 | + xml_obj = read_energyml_xml_str(obj) |
| 183 | + self.energyml_objects.append(xml_obj) |
| 184 | + except: |
| 185 | + try: |
| 186 | + if isinstance(obj, BytesIO): |
| 187 | + obj.seek(0) |
| 188 | + json_obj = read_energyml_json_str(obj, json_version=JSON_VERSION.OSDU_OFFICIAL) |
| 189 | + self.add_file(json_obj) |
| 190 | + except: |
| 191 | + if isinstance(obj, BytesIO): |
| 192 | + obj.seek(0) |
| 193 | + self.add_file(RawFile(path=f"pleaseRenameThisFile_{str(random.random())}.txt", content=obj)) |
| 194 | + elif str(type(obj).__module__).startswith("energyml."): |
| 195 | + # We should test "energyml.(resqml|witsml|prodml|eml|common)" but I didn't to avoid issues if |
| 196 | + # another specific package comes in the future |
| 197 | + self.energyml_objects.append(obj) |
| 198 | + else: |
| 199 | + logging.error(f"unsupported type {str(type(obj))}") |
| 200 | + |
133 | 201 | # EXPORT functions |
134 | 202 |
|
135 | 203 | def gen_opc_content_type(self) -> Types: |
|
0 commit comments