Coverage for oc_meta / lib / file_manager.py: 79%
154 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-07-25 10:39 +0000
1#!/usr/bin/python
3# SPDX-FileCopyrightText: 2022-2026 Arcangelo Massari <arcangelo.massari@unibo.it>
4#
5# SPDX-License-Identifier: ISC
8from __future__ import annotations
10import csv
11import os
12import sys
13from contextlib import contextmanager
14from pathlib import Path
15from time import sleep
16from typing import Callable, Dict, List, Set
17from zipfile import ZIP_DEFLATED, ZipFile
19import orjson
20from _collections_abc import dict_keys
21from bs4 import BeautifulSoup
22from requests import ReadTimeout, get
23from requests.exceptions import ConnectionError
24from scandir_rs import Walk # type: ignore[import-untyped]
26from oc_ocdm.support.support import find_paths, parse_uri
28from oc_meta.lib.cleaner import normalize_spaces
31def find_rdf_file(
32 uri: str,
33 base_dir: str,
34 dir_split: int,
35 items_per_file: int,
36 zip_output: bool = False,
37) -> str:
38 if base_dir and not base_dir.endswith(os.sep):
39 base_dir += os.sep
40 base_iri = parse_uri(uri).base_iri
41 _, file_path = find_paths(uri, base_dir, base_iri, "", dir_split, items_per_file)
42 if zip_output:
43 file_path = os.path.splitext(file_path)[0] + ".zip"
44 return file_path
47def collect_files(
48 root: str,
49 pattern: str = "*.zip",
50 path_filter: Callable[[str], bool] | None = None,
51) -> List[str]:
52 """
53 Directory traversal to collect files matching a pattern.
55 Uses scandir-rs (Rust-based) for fast directory iteration.
57 :param root: Root directory to start traversal
58 :param pattern: Glob pattern for filenames (e.g., '*.zip', 'se.zip')
59 :param path_filter: Optional callable that receives full file path and returns
60 True to include, False to exclude. Example:
61 lambda p: 'prov' not in p # exclude prov directories
62 :returns: List of matching file paths
63 """
64 collected: List[str] = []
65 for dirpath, _, filenames in Walk(root, file_include=[pattern]):
66 for filename in filenames:
67 full_path = os.path.join(root, dirpath, filename)
68 if path_filter is None or path_filter(full_path):
69 collected.append(full_path)
70 return collected
73def collect_zip_files(
74 root: str,
75 only_data: bool = False,
76 only_prov: bool = False,
77) -> List[str]:
78 """
79 Collect ZIP files from a directory tree.
81 :param root: Root directory to start traversal
82 :param only_data: Only include files NOT in paths containing 'prov'
83 :param only_prov: Only include files in paths containing 'prov'
84 :returns: Sorted list of ZIP file paths
86 If both only_data and only_prov are False, all ZIP files are collected.
87 """
88 if only_data and only_prov:
89 return []
91 def path_filter(p: str) -> bool:
92 is_provenance = "prov" in Path(os.path.relpath(p, root)).parts
93 return not is_provenance if only_data else is_provenance
95 active_filter = path_filter if only_data or only_prov else None
96 files = collect_files(root, "*.zip", active_filter)
97 return sorted(files)
100def get_csv_data(filepath: str, clean_data: bool = True) -> List[Dict[str, str]]:
101 if not os.path.splitext(filepath)[1].endswith(".csv"):
102 return list()
103 field_size_changed = False
104 cur_field_size = 128
105 data: List[Dict[str, str]] = list()
106 while True:
107 try:
108 with open(filepath, "r", encoding="utf8") as f:
109 if clean_data:
110 lines = (normalize_spaces(line.replace("\0", "")) for line in f)
111 data = list(csv.DictReader(lines, delimiter=","))
112 else:
113 data = list(csv.DictReader(f, delimiter=","))
114 break
115 except csv.Error:
116 cur_field_size *= 2
117 csv.field_size_limit(cur_field_size)
118 field_size_changed = True
119 if field_size_changed:
120 csv.field_size_limit(128)
121 return data
124def pathoo(path):
125 if not os.path.isdir(os.path.dirname(path)):
126 os.makedirs(os.path.dirname(path))
129def write_csv(
130 path: str,
131 datalist: List[dict],
132 fieldnames: list | dict_keys | None = None,
133 method: str = "w",
134) -> None:
135 if datalist:
136 fieldnames = datalist[0].keys() if fieldnames is None else fieldnames
137 pathoo(path)
138 file_exists = os.path.isfile(path)
139 with open(path, method, newline="", encoding="utf-8") as output_file:
140 dict_writer = csv.DictWriter(
141 f=output_file,
142 fieldnames=fieldnames,
143 delimiter=",",
144 quotechar='"',
145 quoting=csv.QUOTE_NONNUMERIC,
146 )
147 if method == "w" or (method == "a" and not file_exists):
148 dict_writer.writeheader()
149 dict_writer.writerows(datalist)
152def normalize_path(path: str) -> str:
153 normal_path = path.replace("\\", "/").replace("/", os.sep)
154 return normal_path
157def init_cache(cache_filepath: str | None) -> Set[str]:
158 completed = set()
159 if cache_filepath:
160 if not os.path.exists(cache_filepath):
161 pathoo(cache_filepath)
162 else:
163 with open(cache_filepath, "r", encoding="utf-8") as cache_file:
164 completed = {line.rstrip("\n") for line in cache_file}
165 return completed
168@contextmanager
169def suppress_stdout():
170 with open(os.devnull, "w") as devnull: # pragma: no cover
171 old_stdout = sys.stdout
172 sys.stdout = devnull
173 try:
174 yield
175 finally:
176 sys.stdout = old_stdout
179def sort_files(files_to_be_processed: list) -> list:
180 if all(
181 filename.replace(".csv", "").isdigit() for filename in files_to_be_processed
182 ):
183 files_to_be_processed = sorted(
184 files_to_be_processed,
185 key=lambda filename: int(filename.replace(".csv", "")),
186 )
187 elif all(
188 filename.split("_")[-1].replace(".csv", "").isdigit()
189 for filename in files_to_be_processed
190 ):
191 files_to_be_processed = sorted(
192 files_to_be_processed,
193 key=lambda filename: int(filename.split("_")[-1].replace(".csv", "")),
194 )
195 return files_to_be_processed
198def zipdir(path, ziph):
199 for root, _, files in os.walk(path):
200 for file in files:
201 ziph.write(
202 os.path.join(root, file),
203 os.path.relpath(os.path.join(root, file), os.path.join(path, "..")),
204 )
207def zipit(dir_list: list, zip_name: str) -> None:
208 zipf = ZipFile(file=zip_name, mode="w", compression=ZIP_DEFLATED, allowZip64=True)
209 for dir in dir_list:
210 zipdir(dir, zipf)
211 zipf.close()
214def zip_files_in_dir(src_dir: str, dst_dir: str, replace_files: bool = False) -> None:
215 """
216 This method zips files individually in all directories starting from a specified root directory.
217 In other words, this function does not zip the entire folder but individual files
218 while maintaining the folder hierarchy in the specified output directory.
220 :params src_dir: the source directory
221 :type src_dir: str
222 :params dst_dir: the destination directory
223 :type dst_dir: str
224 :params replace_files: True if you want to replace the original unzipped files with their zipped versions. The dafult value is False
225 :type replace_files: bool
226 :returns: None
227 """
228 for dirpath, _, filenames in os.walk(src_dir):
229 for filename in filenames:
230 src_path = os.path.join(dirpath, filename)
231 dst_path = os.path.join(
232 dst_dir, str(Path(src_path).parent).replace(f"{src_dir}{os.sep}", "")
233 )
234 if not os.path.exists(dst_path):
235 os.makedirs(dst_path)
236 _, ext = os.path.splitext(filename)
237 zip_path = os.path.join(dst_path, filename).replace(ext, ".zip")
238 with ZipFile(
239 file=zip_path, mode="w", compression=ZIP_DEFLATED, allowZip64=True
240 ) as zipf:
241 zipf.write(src_path, arcname=filename)
242 if replace_files:
243 os.remove(src_path)
246def unzip_files_in_dir(src_dir: str, dst_dir: str, replace_files: bool = False) -> None:
247 """
248 This method unzips zipped files individually in all directories starting from a specified root directory.
249 In other words, this function does not unzip the entire folder but individual files
250 while maintaining the folder hierarchy in the specified output directory.
252 :params src_dir: the source directory
253 :type src_dir: str
254 :params dst_dir: the destination directory
255 :type dst_dir: str
256 :params replace_files: True if you want to replace the original zipped files with their unzipped versions, defaults to [False]
257 :type replace_files: bool
258 :returns: None
259 """
260 for dirpath, _, filenames in os.walk(src_dir):
261 for filename in filenames:
262 if os.path.splitext(filename)[1] == ".zip":
263 src_path = os.path.join(dirpath, filename)
264 dst_path = os.path.join(
265 dst_dir,
266 str(Path(src_path).parent).replace(f"{src_dir}{os.sep}", ""),
267 )
268 if not os.path.exists(dst_path):
269 os.makedirs(dst_path)
270 with ZipFile(file=os.path.join(dst_path, filename), mode="r") as zipf:
271 zipf.extractall(dst_path)
272 if replace_files:
273 os.remove(src_path)
276def read_zipped_json(filepath: str) -> dict | None:
277 """
278 This method reads a zipped json file.
280 :params filepath: the zipped json file path
281 :type src_dir: str
282 :returns: dict -- It returns the json file as a dictionary
283 """
284 with ZipFile(filepath, "r") as zipf:
285 for filename in zipf.namelist():
286 with zipf.open(filename) as f:
287 json_data = f.read()
288 json_dict = orjson.loads(json_data)
289 return json_dict
292def call_api(
293 url: str, headers: dict[str, str], r_format: str = "json"
294) -> dict | BeautifulSoup | None:
295 tentative = 3
296 while tentative:
297 tentative -= 1
298 try:
299 r = get(url, headers=headers, timeout=30)
300 if r.status_code == 200:
301 r.encoding = "utf-8"
302 if r_format == "json":
303 return orjson.loads(r.text)
304 return BeautifulSoup(r.text, "xml")
305 elif r.status_code == 404:
306 return None
307 except ReadTimeout:
308 pass
309 except ConnectionError:
310 sleep(5)
311 return None