Spaces:
Running
Running
File size: 7,466 Bytes
f29e862 fa9fcab f29e862 31086ae f29e862 31086ae f29e862 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 |
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Dict, List, Set, Union
from graphgen.models import RocksDBCache
class ParallelFileScanner:
def __init__(
self, cache_dir: str, allowed_suffix, rescan: bool = False, max_workers: int = 4
):
self.cache = RocksDBCache(os.path.join(cache_dir, "input_paths.db"))
self.allowed_suffix = set(allowed_suffix) if allowed_suffix else None
self.rescan = rescan
self.max_workers = max_workers
def scan(
self, paths: Union[str, List[str]], recursive: bool = True
) -> Dict[str, Any]:
if isinstance(paths, str):
paths = [paths]
results = {}
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
future_to_path = {}
for p in paths:
if os.path.exists(p):
future = executor.submit(
self._scan_files, Path(p).resolve(), recursive, set()
)
future_to_path[future] = p
for future in as_completed(future_to_path):
path = future_to_path[future]
try:
results[path] = future.result()
except Exception as e:
results[path] = {
"error": str(e),
"files": [],
"dirs": [],
"stats": {},
}
return results
def _scan_files(
self, path: Path, recursive: bool, visited: Set[str]
) -> Dict[str, Any]:
path_str = str(path)
# Avoid cycles due to symlinks
if path_str in visited:
return self._empty_result(path_str)
# cache check
cache_key = f"scan::{path_str}::recursive::{recursive}"
cached = self.cache.get(cache_key)
if cached and not self.rescan:
return cached["data"]
files, dirs = [], []
stats = {"total_size": 0, "file_count": 0, "dir_count": 0, "errors": 0}
try:
path_stat = path.stat()
if path.is_file():
return self._scan_single_file(path, path_str, path_stat)
if path.is_dir():
with os.scandir(path_str) as entries:
for entry in entries:
try:
entry_stat = entry.stat(follow_symlinks=False)
if entry.is_dir():
dirs.append(
{
"path": entry.path,
"name": entry.name,
"mtime": entry_stat.st_mtime,
}
)
stats["dir_count"] += 1
else:
# allowed suffix filter
if not self._is_allowed_file(Path(entry.path)):
continue
files.append(
{
"path": entry.path,
"name": entry.name,
"size": entry_stat.st_size,
"mtime": entry_stat.st_mtime,
}
)
stats["total_size"] += entry_stat.st_size
stats["file_count"] += 1
except OSError:
stats["errors"] += 1
except (PermissionError, FileNotFoundError, OSError) as e:
return {"error": str(e), "files": [], "dirs": [], "stats": stats}
if recursive:
sub_visited = visited | {path_str}
sub_results = self._scan_subdirs(dirs, sub_visited)
for sub_data in sub_results.values():
files.extend(sub_data.get("files", []))
stats["total_size"] += sub_data["stats"].get("total_size", 0)
stats["file_count"] += sub_data["stats"].get("file_count", 0)
result = {"path": path_str, "files": files, "dirs": dirs, "stats": stats}
self._cache_result(cache_key, result, path)
return result
def _scan_single_file(
self, path: Path, path_str: str, stat: os.stat_result
) -> Dict[str, Any]:
"""Scan a single file and return its metadata"""
if not self._is_allowed_file(path):
return self._empty_result(path_str)
return {
"path": path_str,
"files": [
{
"path": path_str,
"name": path.name,
"size": stat.st_size,
"mtime": stat.st_mtime,
}
],
"dirs": [],
"stats": {
"total_size": stat.st_size,
"file_count": 1,
"dir_count": 0,
"errors": 0,
},
}
def _scan_subdirs(self, dir_list: List[Dict], visited: Set[str]) -> Dict[str, Any]:
"""
Parallel scan subdirectories
:param dir_list
:param visited
:return:
"""
results = {}
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
futures = {
executor.submit(self._scan_files, Path(d["path"]), True, visited): d[
"path"
]
for d in dir_list
}
for future in as_completed(futures):
path = futures[future]
try:
results[path] = future.result()
except Exception as e:
results[path] = {
"error": str(e),
"files": [],
"dirs": [],
"stats": {},
}
return results
def _cache_result(self, key: str, result: Dict, path: Path):
"""Cache the scan result"""
self.cache.set(
key,
{
"data": result,
"dir_mtime": path.stat().st_mtime,
"cached_at": time.time(),
},
)
def _is_allowed_file(self, path: Path) -> bool:
"""Check if the file has an allowed suffix"""
if self.allowed_suffix is None:
return True
suffix = path.suffix.lower().lstrip(".")
return suffix in self.allowed_suffix
def invalidate(self, path: str):
"""Invalidate cache for a specific path"""
path = Path(path).resolve()
keys = [k for k in self.cache if k.startswith(f"scan::{path}")]
for k in keys:
self.cache.delete(k)
def close(self):
self.cache.close()
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
@staticmethod
def _empty_result(path: str) -> Dict[str, Any]:
return {
"path": path,
"files": [],
"dirs": [],
"stats": {"total_size": 0, "file_count": 0, "dir_count": 0, "errors": 0},
}
|