Found out that we need to use function_info and updated export.py acc…

…ordingly

nbmodular/cell2func.py

@@ -252,8 +252,8 @@ def __init__(
         self.is_class = False
         self.is_pipeline = False
         self.pipeline_name_or_default = None
-        self.api = True
-        self.keep_original_in_documentation = False
+        # self.api = True
+        # self.keep_original_in_documentation = False
         super().__init__(
             **kwargs,
         )
@@ -1362,10 +1362,10 @@ def __init__(
         **kwargs : dict
             Additional keyword arguments.
 
-        Each boolean parameter can be passed as input or by command line. 
-        The function `set_function_action_and_io_args` will replace the None 
-        default values with either True or False depending on whether 
-        `--parameter` or `--not-parameter` was passed in the command line, 
+        Each boolean parameter can be passed as input or by command line.
+        The function `set_function_action_and_io_args` will replace the None
+        default values with either True or False depending on whether
+        `--parameter` or `--not-parameter` was passed in the command line,
         and depending on `default_parameter` if none of these were passed.
 
         Examples
@@ -1406,7 +1406,7 @@ def __init__(
         self.code_cells_path = Path(code_cells_path)
         self.code_cells_path.mkdir(parents=True, exist_ok=True)
 
-        self.api = api # do we need this?
+        self.api = api  # do we need this?
         self.keep_original_in_documentation = keep_original_in_documentation
         self.restrict_inputs = restrict_inputs
         self.current_function = Bunch()
@@ -1547,7 +1547,9 @@ def __init__(
 
         self.default_restrict_inputs = self.restrict_inputs
         self.default_api = self.api
-        self.default_keep_original_in_documentation = self.keep_original_in_documentation
+        self.default_keep_original_in_documentation = (
+            self.keep_original_in_documentation
+        )
 
         self.parser = argparse.ArgumentParser(
             description="Arguments to `function` magic cell."
@@ -1950,7 +1952,7 @@ def set_value(self, attr, value, convert=True):
                 f"Cell processor has no attribute {attr}. Existing attributes are:\n{list_of_attributes}"
             )
 
-    def set_api (self, value):
+    def set_api(self, value):
         self.api = value
         self.default_api = value
 
@@ -3504,8 +3506,8 @@ def set(self, line):
 
     @line_magic
     def keep_original(self, line):
-        self.processor.set_api (False)
-    
+        self.processor.set_api(False)
+
 
 # %% ../nbs/cell2func.ipynb 74
 def load_ipython_extension(ipython):

nbmodular/export.py

@@ -36,6 +36,7 @@
 import sys
 
 # 3rd party
+from matplotlib.pylab import f
 from sklearn.utils import Bunch
 from nbdev.processors import Processor, NBProcessor
 from nbdev.export import nb_export
@@ -315,8 +316,13 @@ def cell(self, cell):
                     is_class=command == "class",
                 )
         elif len(source_lines) > 0 and source_lines[0].strip().startswith("%"):
+            # We process here those cells that start with a line magic command.
+            # Those are not processed when calling self.cell_processor.process_function_call
+            # And they *usually* affect the whole cell_processor, not just one of its functions
+            # Therefore, they are used here to update the state of self.cell_processor
+            # By adding / modifying attributes modified by these line magics.
             line = source_lines[0].strip()[1:]
-            #source = "\n".join(source_lines[1:])
+            # source = "\n".join(source_lines[1:])
             command, remaining_line = line.split()
             # function_name, kwargs = self.cell_processor.parse_signature(remaining_line)
             if command == "keep_original":
@@ -487,22 +493,31 @@ def cell(self, cell):
                     if is_test
                     else self.nb_magic_processor.cell_processor.code_cells
                 )
+                function_info = (
+                    self.nb_magic_processor.cell_processor.function_info
+                    if is_test
+                    else self.nb_magic_processor.cell_processor.function_info
+                )
                 if function_name not in code_cells:
                     raise RuntimeError(
                         f"Function {function_name} not found in code_cells dictionary with keys {code_cells.keys()}"
                     )
+                function_info = function_info[function_name]
                 code_cells = code_cells[function_name]
                 if len(code_cells) <= idx:
                     raise RuntimeError(
                         f"Function {function_name} has {len(code_cells)} cells, which is lower than index {idx}."
                     )
                 code_cell = code_cells[idx]
                 self.logger.debug("code:")
+                # should we use code_cell for storing valid flag, or function_info?
+                # (this depends on whether the valid flag needs to be saved as part
+                # of the pickle file for code_cells)
                 self.logger.debug(f"{code_cell.code}valid: {code_cell.valid}")
-                # keep_original_in_documentation = (
-                #     code_cell.keep_original_in_documentation
-                # )
-                if code_cell.valid and code_cell.api:
+                keep_original_in_documentation = (
+                    function_info.keep_original_in_documentation
+                )
+                if code_cell.valid and function_info.api:
                     source = code_cell.code
                     to_export = True
             elif line.startswith("%%include") or line.startswith("%%class"):
@@ -523,11 +538,11 @@ def cell(self, cell):
                     self.cells.append(new_cell)
                     cell_type = "code"
             else:
-                # if keep_original_in_documentation:
-                #     doc_source = cell.source
-                # else:
-                #     doc_source = source  # doc_source does not include first line with %% (? to think about)
-                doc_source = source  # doc_source does not include first line with %% (? to think about)
+                if keep_original_in_documentation:
+                    doc_source = cell.source
+                else:
+                    doc_source = source  # doc_source does not include first line with %% (? to think about)
+                # doc_source = source  # doc_source does not include first line with %% (? to think about)
             if is_test:
                 doc_source = transform_test_source_for_docs(
                     code_cell.code, idx, self.tab_size
@@ -540,7 +555,7 @@ def cell(self, cell):
             and source_lines[0].strip().startswith("%")
             and "--keep" in source_lines[0].strip()
         ):
-            self.nb_magic_processor.cell_processor.code_cells.append (cell)
+            self.nb_magic_processor.cell_processor.code_cells.append(cell)
 
         self.cell_types.append(cell_type)