Update stats.py

aws_doc_sdk_examples_tools/stats.py

@@ -1,28 +1,55 @@
-from typing import List
+from typing import List, Dict, Any
 from pathlib import Path
 from pprint import pformat
-
+import logging
 from .doc_gen import DocGen
 
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-def main(roots: List[str]):
-    base = DocGen.empty()
-    for root in roots:
+def process_root(root: str, base: DocGen) -> Dict[str, Any]:
+    """
+    Process a single root directory & collect example stats.
+
+    Args:
+        root (str): The root dir to process.
+        base (DocGen): Instance of DocGen class.
+
+    Returns:
+        dict: Stats for the given root.
+    """
+    try:
         docgen_root = Path(root)
+        if not docgen_root.exists():
+            logging.warning(f"Root directory {docgen_root} does not exist.")
+            return {}
+
         doc_gen = base.clone().for_root(docgen_root)
         doc_gen.collect_snippets()
-        print(f"Root	{docgen_root.name}")
+
         stats = doc_gen.stats()
-        print(f"SDKs	{stats['sdks']}")
-        print(f"Services	{stats['services']}")
-        print(f"Examples	{stats['examples']}")
-        print(f"Version	{stats['versions']}")
-        print(f"Snippets	{stats['snippets']}")
-        genai = pformat(dict(stats["genai"]))
-        print(f"GenAI	{genai}")
+        stats['root'] = docgen_root.name
+        return stats
+    except Exception as e:
+        logging.error(f"Error processing root {root}: {e}")
+        return {}
+
+def collect_stats(roots: List[str]) -> List[Dict[str, Any]]:
+    """
+    Collects stats for a list of root directories.
+
+    Args:
+        roots (List[str]): A list of root directories as strings.
 
+    Returns:
+        List[Dict[str, Any]]: A list of dictionaries with collected stats for each root.
+    """
+    base = DocGen.empty()
+    all_stats = []
+
+    for root in roots:
+        stats = process_root(root, base)
+        if stats:
+            all_stats.append(stats)
 
-if __name__ == "__main__":
-    from sys import argv
+    return all_stats
 
-    main(argv[1:])