Coverage for oc_meta / run / count / triples.py: 98%
163 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-07-25 10:39 +0000
1# SPDX-FileCopyrightText: 2026 Arcangelo Massari <arcangelo.massari@unibo.it>
2#
3# SPDX-License-Identifier: ISC
5"""Count RDF triples or quads in compressed or uncompressed files."""
7from __future__ import annotations
9import argparse
10import fnmatch
11import gzip
12import multiprocessing
13import os
14import zipfile
15from collections.abc import Iterable, Mapping, Sequence
16from functools import partial
17from pathlib import Path
18from typing import TextIO
20import orjson
21from rich_argparse import RichHelpFormatter
23from oc_meta.lib.console import create_progress
24from oc_meta.lib.file_manager import collect_files
26QUAD_FORMATS = {"nquads"}
27LINE_BASED_FORMATS = {"nquads", "nt"}
28JSONLD_SPECIAL_KEYS = {
29 "@id",
30 "@context",
31 "@graph",
32 "@list",
33 "@set",
34 "@language",
35 "@value",
36}
39def _count_jsonld_value(value: object) -> int:
40 """Count triples generated by a JSON-LD property value."""
41 if isinstance(value, list):
42 return sum(_count_jsonld_value(item) for item in value)
43 if isinstance(value, dict):
44 if "@list" in value:
45 return _count_jsonld_value(value["@list"])
46 if "@set" in value:
47 return _count_jsonld_value(value["@set"])
48 if "@value" in value:
49 return 1
50 if "@id" in value and len(value) == 1:
51 return 1
52 return 1 + _count_jsonld_object(value)
53 return 1
56def _count_jsonld_object(obj: Mapping[str, object]) -> int:
57 """Count triples for a single JSON-LD object."""
58 count = 0
59 for key, value in obj.items():
60 if key in JSONLD_SPECIAL_KEYS:
61 continue
62 if key == "@type":
63 if isinstance(value, list):
64 count += len(value)
65 else:
66 count += 1
67 else:
68 count += _count_jsonld_value(value)
69 return count
72def _count_jsonld_triples(data: Mapping[str, object] | Sequence[object]) -> int:
73 """Count triples in a JSON-LD document without RDFLib parsing."""
74 if isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
75 total = 0
76 for obj in data:
77 if isinstance(obj, dict):
78 if "@graph" in obj:
79 graph = obj["@graph"]
80 if isinstance(graph, list):
81 total += sum(
82 _count_jsonld_object(item)
83 for item in graph
84 if isinstance(item, dict)
85 )
86 else:
87 total += _count_jsonld_object(obj)
88 return total
89 if isinstance(data, Mapping):
90 if "@graph" in data:
91 graph = data["@graph"]
92 if isinstance(graph, list):
93 return sum(
94 _count_jsonld_object(obj) for obj in graph if isinstance(obj, dict)
95 )
96 return _count_jsonld_object(data)
97 return 0
100def parse_args() -> argparse.Namespace: # pragma: no cover
101 parser = argparse.ArgumentParser(
102 description="Count RDF triples or quads in compressed or uncompressed files.",
103 formatter_class=RichHelpFormatter,
104 )
105 parser.add_argument(
106 "directory",
107 type=Path,
108 help="Directory containing the RDF files.",
109 )
110 parser.add_argument(
111 "--pattern",
112 default="*.nq.gz",
113 help="Glob pattern for locating files (default: '*.nq.gz').",
114 )
115 parser.add_argument(
116 "--format",
117 default="nquads",
118 choices=["nquads", "nt", "json-ld"],
119 help="RDF format of the input files (default: nquads).",
120 )
121 parser.add_argument(
122 "--recursive",
123 action="store_true",
124 help="Search recursively under the provided directory.",
125 )
126 parser.add_argument(
127 "--prov-only",
128 action="store_true",
129 help="Count only files in 'prov' subdirectories.",
130 )
131 parser.add_argument(
132 "--data-only",
133 action="store_true",
134 help="Count only files not in 'prov' subdirectories.",
135 )
136 parser.add_argument(
137 "--workers",
138 type=int,
139 default=None,
140 help="Number of parallel workers (default: CPU count).",
141 )
142 parser.add_argument(
143 "--show-per-file",
144 action="store_true",
145 help="Print the count for each processed file.",
146 )
147 parser.add_argument(
148 "--keep-going",
149 action="store_true",
150 help="Continue processing even if errors occur.",
151 )
152 return parser.parse_args()
155def discover_files(
156 directory: Path,
157 pattern: str,
158 recursive: bool,
159 prov_only: bool,
160 data_only: bool,
161) -> list[Path]:
162 path = directory.expanduser().resolve()
163 if not path.is_dir():
164 raise ValueError(f"'{path}' does not exist or is not a directory.")
166 root_str = str(path)
168 if recursive:
170 def path_filter(p: str) -> bool:
171 is_prov = "/prov" in p or p.endswith("/prov")
172 if prov_only and not is_prov:
173 return False
174 if data_only and is_prov:
175 return False
176 return True
178 str_files = collect_files(root_str, pattern, path_filter)
179 return sorted(Path(f) for f in str_files)
181 is_prov = path.name == "prov"
182 if (prov_only and not is_prov) or (data_only and is_prov):
183 return []
185 files: list[Path] = []
186 for entry in os.scandir(root_str):
187 if entry.is_file() and fnmatch.fnmatch(entry.name, pattern):
188 files.append(Path(entry.path))
189 return sorted(files)
192def _count_lines_binary(file_obj: Iterable[bytes]) -> int:
193 count = 0
194 for line_num, line in enumerate(file_obj, 1):
195 stripped = line.strip()
196 if not stripped or stripped.startswith(b"#"):
197 continue
198 if not stripped.endswith(b"."):
199 raise ValueError(f"line {line_num}: statement does not end with '.'")
200 count += 1
201 return count
204def _count_lines_text(file_obj: TextIO) -> int:
205 count = 0
206 for line_num, line in enumerate(file_obj, 1):
207 stripped = line.strip()
208 if not stripped or stripped.startswith("#"):
209 continue
210 if not stripped.endswith("."):
211 raise ValueError(f"line {line_num}: statement does not end with '.'")
212 count += 1
213 return count
216def count_in_file(file_path: Path, rdf_format: str) -> tuple[str, int, str | None]:
217 try:
218 suffix = file_path.suffix.lower()
219 use_line_count = rdf_format in LINE_BASED_FORMATS
221 if suffix == ".zip":
222 with zipfile.ZipFile(file_path, "r") as z:
223 inner_name = z.namelist()[0]
224 with z.open(inner_name) as f:
225 if use_line_count:
226 return str(file_path), _count_lines_binary(f), None
227 content = f.read().decode("utf-8")
228 data = orjson.loads(content)
229 return str(file_path), _count_jsonld_triples(data), None
230 elif suffix == ".gz":
231 if use_line_count:
232 with gzip.open(file_path, "rb") as f:
233 return str(file_path), _count_lines_binary(f), None
234 with gzip.open(file_path, "rb") as f:
235 data = orjson.loads(f.read())
236 return str(file_path), _count_jsonld_triples(data), None
237 else:
238 if use_line_count:
239 with open(file_path, "r", encoding="utf-8") as f:
240 return str(file_path), _count_lines_text(f), None
241 with open(file_path, "rb") as f:
242 data = orjson.loads(f.read())
243 return str(file_path), _count_jsonld_triples(data), None
244 except Exception as exc:
245 return str(file_path), 0, str(exc)
248def process_files(
249 files: list[Path],
250 rdf_format: str,
251 max_workers: int | None,
252 show_per_file: bool,
253 keep_going: bool,
254 unit_name: str,
255) -> tuple[int, list[tuple[str, str]]]:
256 workers = max_workers or multiprocessing.cpu_count()
257 if workers < 1:
258 workers = 1
260 total_count = 0
261 results: list[tuple[str, int]] = []
262 failures: list[tuple[str, str]] = []
263 chunksize = max(1, len(files) // (workers * 4))
265 worker_fn = partial(count_in_file, rdf_format=rdf_format)
267 with create_progress() as progress:
268 task = progress.add_task(f"Counting {unit_name}", total=len(files))
270 # Use forkserver to avoid deadlocks when forking in a multi-threaded environment
271 ctx = multiprocessing.get_context("forkserver")
272 with ctx.Pool(processes=workers) as pool:
273 for file_path, count, error in pool.imap_unordered(
274 worker_fn, files, chunksize=chunksize
275 ):
276 if error:
277 failures.append((file_path, error))
278 if not keep_going:
279 progress.console.print(
280 f"[red]Error processing {file_path}: {error}[/red]"
281 )
282 progress.advance(task)
283 pool.terminate()
284 break
285 progress.console.print(
286 f"[yellow]Error processing {file_path}: {error} (continuing)[/yellow]"
287 )
288 else:
289 total_count += count
290 if show_per_file:
291 results.append((file_path, count))
293 progress.advance(task)
295 if show_per_file and results:
296 width = max(len(path) for path, _ in results)
297 for path, count in sorted(results):
298 print(f"{path.ljust(width)} : {count}")
300 return total_count, failures
303def main() -> None: # pragma: no cover
304 args = parse_args()
306 if args.prov_only and args.data_only:
307 print("Error: --prov-only and --data-only are mutually exclusive.")
308 return
310 try:
311 files = discover_files(
312 args.directory,
313 args.pattern,
314 args.recursive,
315 args.prov_only,
316 args.data_only,
317 )
318 except ValueError as exc:
319 print(f"Error: {exc}")
320 return
322 if not files:
323 print("No files found matching the provided pattern.")
324 return
326 unit_name = "quads" if args.format in QUAD_FORMATS else "triples"
328 total, failures = process_files(
329 files,
330 args.format,
331 args.workers,
332 args.show_per_file,
333 args.keep_going,
334 unit_name,
335 )
337 print(f"Total {unit_name}: {total}")
339 if failures:
340 print(f"\nFiles with errors ({len(failures)}):")
341 for path, error in failures:
342 print(f" {path}: {error}")
345if __name__ == "__main__": # pragma: no cover
346 main()