feat: add allow_suffix to limit input files

Author: ChenZiHong-Gavin

graphgen/graphgen.py

@@ -96,7 +96,7 @@ async def read(self, read_config: Dict):
         """
         read files from input sources
         """
-        data = read_files(read_config["input_file"], self.working_dir)
+        data = read_files(**read_config, cache_dir=self.working_dir)
         if len(data) == 0:
             logger.warning("No data to process")
             return

graphgen/operators/read/read_files.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from graphgen.models import (
     CSVReader,
@@ -34,26 +34,49 @@ def _build_reader(suffix: str, cache_dir: str | None):
     return _MAPPING[suffix]()
 
 
-def read_files(file_path: str, cache_dir: str | None = None) -> list[dict]:
-    path = Path(file_path).expanduser()
+def read_files(
+    input_file: str,
+    allowed_suffix: Optional[List[str]] = None,
+    cache_dir: Optional[str] = None,
+) -> list[dict]:
+    path = Path(input_file).expanduser()
     if not path.exists():
-        raise FileNotFoundError(f"input_path not found: {file_path}")
+        raise FileNotFoundError(f"input_path not found: {input_file}")
 
+    if allowed_suffix is None:
+        support_suffix = set(_MAPPING.keys())
+    else:
+        support_suffix = {s.lower().lstrip(".") for s in allowed_suffix}
+
+    # single file
     if path.is_file():
-        suffix = path.suffix.lstrip(".")
+        suffix = path.suffix.lstrip(".").lower()
+        if suffix not in support_suffix:
+            logger.warning(
+                "Skip file %s (suffix '%s' not in allowed_suffix %s)",
+                path,
+                suffix,
+                support_suffix,
+            )
+            return []
         reader = _build_reader(suffix, cache_dir)
         return reader.read(str(path))
 
-    support_suffix = set(_MAPPING.keys())
+    # folder
     files_to_read = [
         p for p in path.rglob("*") if p.suffix.lstrip(".").lower() in support_suffix
     ]
-    logger.info("Found %d file(s) under folder %s", len(files_to_read), file_path)
+    logger.info(
+        "Found %d eligible file(s) under folder %s (allowed_suffix=%s)",
+        len(files_to_read),
+        input_file,
+        support_suffix,
+    )
 
     all_docs: List[Dict[str, Any]] = []
     for p in files_to_read:
         try:
-            suffix = p.suffix.lstrip(".")
+            suffix = p.suffix.lstrip(".").lower()
             reader = _build_reader(suffix, cache_dir)
             all_docs.extend(reader.read(str(p)))
         except Exception as e:  # pylint: disable=broad-except