Export imported modules

Closes #18 (closed)

plotid/publish.py

@@ -8,6 +8,7 @@ the plot is based on. Additionally, the script that produced the plot will be
 copied to the destination directory.
 """
 
+import ast
 import os
 import shutil
 import sys
@@ -239,6 +240,55 @@ class PublishOptions:
                 os.path.join(destination, final_file_path),
             )
 
+    def export_imports(self, file: str) -> None:
+        """Export all imported modules of a python script to file."""
+        with open(file, "r", encoding="utf-8") as source:
+            tree = ast.parse(source.read())
+
+        analyzer = Analyzer()
+        analyzer.visit(tree)
+        analyzer.report(self.dst_path)
+
+
+class Analyzer(ast.NodeVisitor):
+    """Visit and analyze nodes of Abstract Syntax Trees (AST)."""
+
+    def __init__(self) -> None:
+        self.stats: dict[str, list[str]] = {
+            "import": [],
+            "from_module": [],
+            "from": [],
+        }
+
+    def visit_Import(self, node: ast.Import) -> None:
+        """Get modules that are imported with the 'import' statement."""
+        for alias in node.names:
+            self.stats["import"].append(alias.name)
+        self.generic_visit(node)
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        """Get modules that are imported with the 'from X import Y' statement."""
+        self.stats["from"].append(str(node.module))
+        self.generic_visit(node)
+
+    def report(self, dst_dir: str) -> None:
+        """Create summary of imported modules."""
+        # Save the first part of import statement since it references the installed
+        # module.
+        imports_as_set = {module.split(".", 1)[0] for module in self.stats["import"]}
+
+        imports_as_set.update(
+            # Add modules imported with "from X import Y".
+            {module.split(".", 1)[0] for module in self.stats["from"]}
+        )
+
+        output_file = os.path.join(dst_dir, "required_imports.txt")
+        # Write every item of the set to one line.
+        with open(output_file, "w", encoding="utf-8") as output:
+            for item in imports_as_set:
+                output.write(f"{item}\n")
+            output.close()
+
 
 kwargs_types_publish = TypedDict(
     "kwargs_types_publish",
@@ -288,3 +338,4 @@ def publish(
     publish_container = PublishOptions(figs_and_ids, src_datapath, dst_path, **kwargs)
     publish_container.validate_input()
     publish_container.export()
+    publish_container.export_imports(sys.argv[0])