Added a CLI flag for __init__.py generation and improved package imports

Author: ThunderStruct

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## 0.1.4 (2024-02-22)
+
+[Compare the full difference.](https://github.com/ThunderStruct/nbrefactor/compare/0.1.3...0.1.4)
+
+### Added
+
+- Added new CLI flag `-i/--init-files` to optionally generate `__init__.py` files in package directories
+- Improved package imports by properly exposing modules in `__init__.py` files
+
+### Fixed
+
+- Fixed relative import issues by properly handling package hierarchy with `__init__.py` files
+
 ## 0.1.3 (2025-02-18)
 
 [Compare the full difference.](https://github.com/ThunderStruct/nbrefactor/compare/0.1.2...0.1.3)

TODO.md

@@ -1,4 +1,3 @@
-
 # Todo
 
 - 🔴 High priority _(i.e. critical and urgent features / fixes)_
@@ -21,7 +20,7 @@
 
 ## Minor
 
-- [ ] 🔴 Add a CLI flag to instantiate modules (add `__init__.py` to sub-directories in the writer phase).
+- [x] 🔴 Add a CLI flag to instantiate modules (add `__init__.py` to sub-directories in the writer phase).
 - [ ] 🟠 Add unit testing for the [CDA](src/nbrefactor/processor/cda.py), [Parser](src/nbrefactor/processor/parser.py), and [CLI](src/nbrefactor/cli.py). Integrate it into the CI workflow.
 
 

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name='nbrefactor',
-    version='0.1.3',
+    version='0.1.4',
     author='Mohamed Shahawy',
     author_email='[email protected]',
     description='An automation tool to refactor Jupyter Notebooks to Python modules, with code dependency analysis.',

src/nbrefactor/__init__.py

@@ -1,5 +1,5 @@
 # Current
-__version__ = '0.1.3'
+__version__ = '0.1.4'
 
 # CLI
 from .cli import main

src/nbrefactor/cli.py

@@ -17,7 +17,8 @@ def main():
     # refactoring
     root_node = process_notebook(notebook_path=args.notebook_path, 
                                  output_path=args.output_path, 
-                                 root_package=args.root_package)
+                                 root_package=args.root_package,
+                                 generate_init=args.generate_init)
 
     # plotting (if applicable)
     if args.generate_plot:

src/nbrefactor/fileops/writer.py

@@ -6,7 +6,7 @@
 from ..datastructs import ParsedCodeCell
 
 
-def write_modules(node, output_path, pre_write_hook=None):
+def write_modules(node, output_path, pre_write_hook=None, generate_init=False):
     """
     Recursively writes the modules and their contents given a \
         :class:`~nbrefactor.datastructs.ModuleNode` tree hierarchy.
@@ -16,6 +16,8 @@ def write_modules(node, output_path, pre_write_hook=None):
         output_path (str): the path where the Python files should be written.
         write_hook (callable, optional): a hook function that takes \
             (content, node) and returns the modified content.
+        generate_init (bool, optional): whether to generate __init__.py files \
+            in package directories. Defaults to False.
     """
 
     if (not node.children and node.has_code_cells()) \
@@ -35,6 +37,26 @@ def write_modules(node, output_path, pre_write_hook=None):
         module_path = os.path.join(output_path, node.name)
         ensure_dir(module_path)
 
+        # create `__init__.py` file for package if requested
+        if generate_init and node.name != '.':  # skip root directory
+            init_path = os.path.join(module_path, '__init__.py')
+            with open(init_path, 'w') as f:
+                # write package docstring
+                f.write('"""Package initialization."""\n\n')
+                
+                # if this node has code cells, expose them at package level
+                if node.has_code_cells():
+                    module_name = os.path.splitext(node.name)[0]
+                    # noqa to supress import warnings (refer to PEP 8)
+                    f.write(f'from .{module_name} import *  # noqa\n')
+                
+                # expose all child modules
+                for child in node.children.values():
+                    if not child.ignore_package and not child.ignore_module:
+                        child_name = os.path.splitext(child.name)[0]
+                        # noqa to supress import warnings
+                        f.write(f'from .{child_name} import *  # noqa\n')
+
         # node could potentially have code at package-level, so...
         if node.has_code_cells():
             __write_module_node(node, module_path, 
@@ -43,7 +65,9 @@ def write_modules(node, output_path, pre_write_hook=None):
 
         # recursively write child nodes
         for child in node.children.values():
-            write_modules(child, module_path, pre_write_hook=pre_write_hook)
+            write_modules(child, module_path, 
+                         pre_write_hook=pre_write_hook,
+                         generate_init=generate_init)
 
 
 def __write_module_node(node, output_path, 

src/nbrefactor/processor/processor.py

@@ -13,7 +13,7 @@
 
 
 def process_notebook(notebook_path, output_path, 
-                     root_package='.', pre_write_hook=None):
+                     root_package='.', pre_write_hook=None, generate_init=False):
     """
     The Notebook-refactoring entry point. This function:
 
@@ -37,6 +37,8 @@ def process_notebook(notebook_path, output_path,
                 `output_path`).
         write_hook (callable, optional): a hook function that takes (content, \
             node) and returns the modified content.
+        generate_init (bool, optional): whether to generate __init__.py files \
+            in package directories. Defaults to False.
     
     Returns:
         :class:`~nbrefactor.datastructs.ModuleNode`: the root node of the \
@@ -165,7 +167,8 @@ def write_hook_wrapper(content, node):
         return content
 
     write_modules(root, output_path, 
-                  pre_write_hook=write_hook_wrapper)
+                  pre_write_hook=write_hook_wrapper,
+                  generate_init=generate_init)
 
     Logger.log((
             f'Successfully wrote ({file_counter}) files to "{output_path}"!\n'

src/nbrefactor/utils/argument_parsing.py

@@ -21,6 +21,8 @@ def parse_args():
             Defaults to `False`.
         -pf, --plot-format (str): format of the plot (e.g., "pdf", "png"). \
             Defaults to "pdf".
+        -i, --init-files: generate __init__.py files in package directories. \
+            Defaults to `False`.
     
     Returns:
         argparse.Namespace: parsed arguments from the command line.
@@ -65,4 +67,9 @@ def parse_args():
                         help=('Format of the plot (e.g., "pdf", "png"). '
                               'Defaults to "pdf".'))
 
+    parser.add_argument('-i', '--init-files',
+                        action='store_true',
+                        dest='generate_init',
+                        help='Generate __init__.py files in package directories.')
+
     return parser.parse_args()