utils / acb : added commandline argument

Just in case we may want to process only one file

Author: beachdweller

utils/add_colab_main_buttons.py

@@ -1,3 +1,4 @@
+import argparse
 import copy
 import functools
 import json
@@ -8,7 +9,7 @@
 import urllib.parse as up
 
 
-from typing import Dict, Tuple
+from typing import Dict, List, Tuple
 
 
 import bs4
@@ -19,16 +20,23 @@
 
 
 def main(argv):
-    if len(argv) > 1:
-        full_path = pathlib.Path(argv[1]).absolute()
-        if full_path.is_file():
-            proc_file(full_path)
-        elif full_path.is_dir():
-            proc_dir(full_path)
-        else:
-            raise NotImplementedError(full_path)
+    parsed = parse_argv(argv)
+
+    if parsed.file != pathlib.Path('None') and parsed.file.is_file():
+        proc_file(parsed.file)
     else:
-        proc_dir()
+        proc_dir(parsed.directory)
+
+
+def parse_argv(argv:List[str]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+
+    # Just in case to process only one file
+    parser.add_argument('-f', "--file", required=False, type=pathlib.Path, default='None')
+    # Just in case to process only one folder
+    parser.add_argument('-d', "--directory", required=False, type=pathlib.Path, default=get_proj_root())
+
+    return parser.parse_args(argv[1:])
 
 
 def proc_dir(root:str=None):
@@ -131,11 +139,10 @@ def has_button_img(cell:Dict) -> bool:
 
 @functools.lru_cache()
 def get_proj_root() -> pathlib.Path:
-    result = pathlib.Path(__file__).parent.parent.absolute()
-    assert result.exists(), result
-    assert result.is_dir()
-    assert (result / ".gitignore").exists(), (result, result.glob("*"))
-    return result
+    '''
+    Cache version of rsc.get_proj_root()
+    '''
+    return rsc.get_proj_root()
 
 
 def get_rel_path(full_path:str) -> str:

utils/tests/test_add_colab_main_buttons.py

@@ -107,5 +107,23 @@ def test_has_button_img__code_cell(self):
         self.assertFalse(result)
 
 
+class TestAddColabMainButtonsParseArgv(unittest.TestCase):
+    def test_parse_argv__no_args(self):
+        result = acb.parse_argv([])
+        self.assertEqual(result.file, pathlib.Path('None'))
+        self.assertEqual(result.directory, acb.get_proj_root())
+
+    def test_parse_argv__one_file_arg(self):
+        filename = "test.ipynb"
+        result = acb.parse_argv(['dummy.py', "-f", "test.ipynb"])
+        self.assertEqual(result.file, pathlib.Path(filename))
+
+    def test_parse_argv__one_dir_arg(self):
+        dirname = "test_folder"
+        result = acb.parse_argv(['dummy.py', "-d", dirname])
+        self.assertEqual(result.file, pathlib.Path('None'))
+        self.assertEqual(result.directory, pathlib.Path(dirname))
+
+
 if "__main__" == __name__:
     unittest.main()