feat: implement AddImportToFile tranformation

geny/core/filesystem/directories.py

@@ -64,7 +64,7 @@ def path(self, parent: Path = None) -> Path:
         if parent is None:
             return Path(self.name)
 
-        if type(parent) is str:
+        if isinstance(parent, str):
             parent = Path(parent)
 
         return parent / self.name

geny/core/filesystem/transformations.py

@@ -1,4 +1,5 @@
 # core:filesystem:transformations
+import re
 from typing import Protocol
 from pathlib import Path
 
@@ -62,6 +63,36 @@ def run(self):
             pass
 
 
+class AddImportToFile(AddLineToFile):
+    def run(self):
+        if self.statement is None or self.statement == "":
+            return
+
+        try:
+            with open(self.target, mode="r+") as f:
+                lines = f.readlines() or []
+
+                header_border = 0
+                for idx, line in enumerate(lines):
+                    if self.prevent_duplicates and line.startswith(self.statement):
+                        return
+
+                    # Matches against comments and import statements
+                    if re.match('(^\s+$|^#.*|^"""|^import [a-zA-Z_.]+( as [a-zA-Z_.]+)?(\s+#.*)?$)|(^from [a-zA-Z_.]+ import [a-zA-Z_.]+( as [a-zA-Z_.]+)?(\s+#.*)?$)', line):
+                        header_border = idx
+                    else:
+                        # Found a 'program line'!!
+                        break
+
+                # Import statement should go before first 'program line'
+                lines.insert(header_border, f"{self.statement}\n\n")
+
+                f.truncate(0)
+                f.writelines(lines)
+        except (FileNotFoundError, OSError) as _:
+            pass
+
+
 class RemoveLineFromFile(FileTransformation):
     def __init__(self, target: Path, statement: str):
         self.target = target
@@ -84,3 +115,4 @@ def run(self):
                 f.writelines(lines)
         except (FileNotFoundError, OSError) as _:
             pass
+

pyproject.toml

@@ -6,7 +6,7 @@ maintainers = [{ name = "Leo Neto", email = "[email protected]" }]
 requires-python = ">=3.6"
 license = { file = "LICENSE" }
 readme = "README.md"
-version = "0.1.6"
+version = "0.1.7"
 keywords = [
   "automation",
   "files",

tests/core/filesystem/test_transformations.py

@@ -6,6 +6,7 @@
     DeleteFile,
     MoveFile,
     AddLineToFile,
+    AddImportToFile,
     RemoveLineFromFile,
 )
 
@@ -25,6 +26,45 @@ def test_add_line(self):
             AddLineToFile(file, "line_3").run()
             self.assertEqual("line_1\nline_2\nline_3\n", file.read_text("utf-8"))
 
+    def test_add_import_line(self):
+        file = Path("new.scratch")
+
+        with runner.isolated_filesystem():
+            file.write_text("""
+# example
+import sys
+import math
+import random
+import numpy as np # used for matrix stuff
+
+
+from django.db import models # TODO: work with models
+# from pathlib import Path
+
+# TODO: implement
+def main():
+    return "greetings"
+
+if __name__ == "__main__":
+    from pprint import pprint as p
+    p(main())
+""")
+
+            stmt = "from os import listdir as ls"
+            AddImportToFile(file, stmt).run()
+            self.assertIn("from os import listdir as ls", file.read_text("utf-8"))
+
+            with open(file, 'r') as f:
+                lines = f.readlines()
+
+                # import statement
+                self.assertIn(stmt, lines[11])
+
+                # surrounding lines (before and after)
+                self.assertIn("# from pathlib import Path", lines[9])
+                self.assertIn("# TODO: implement", lines[13])
+
+
     def test_remove_line(self):
         file = Path("new.scratch")