add figure annotation

Author: giannisdoukas

examples/intro.ipynb

@@ -15,7 +15,7 @@
 from nbformat.notebooknode import NotebookNode  # type: ignore
 
 from .iotypes import CWLFilePathInput, CWLBooleanInput, CWLIntInput, CWLStringInput, CWLFilePathOutput, \
-    CWLDumpableFile, CWLDumpableBinaryFile, CWLDumpable, CWLPNGPlot
+    CWLDumpableFile, CWLDumpableBinaryFile, CWLDumpable, CWLPNGPlot, CWLPNGFigure
 from .requirements_manager import RequirementsManager
 
 with open(os.sep.join([os.path.abspath(os.path.dirname(__file__)), 'templates', 'template.dockerfile'])) as f:
@@ -62,14 +62,19 @@ class AnnotatedVariablesExtractor(ast.NodeTransformer):
 
     dumpable_mapper = {
         (CWLDumpableFile.__name__,): (
-            "with open('{var_name}', 'w') as f:\n\tf.write({var_name})", lambda node: node.target.id
+            (None, "with open('{var_name}', 'w') as f:\n\tf.write({var_name})",),
+            lambda node: node.target.id
         ),
         (CWLDumpableBinaryFile.__name__,): (
-            "with open('{var_name}', 'wb') as f:\n\tf.write({var_name})", lambda node: node.target.id
+            (None, "with open('{var_name}', 'wb') as f:\n\tf.write({var_name})"),
+            lambda node: node.target.id
         ),
         (CWLDumpable.__name__, CWLDumpable.dump.__name__): None,
         (CWLPNGPlot.__name__,): (
-            'import matplotlib.pyplot as plt\nplt.savefig("{var_name}.png")',
+            (None, '{var_name}[-1].figure.savefig("{var_name}.png")'),
+            lambda node: str(node.target.id) + '.png'),
+        (CWLPNGFigure.__name__,): (
+            ('import matplotlib.pyplot as plt\nplt.figure()', '{var_name}[-1].figure.savefig("{var_name}.png")'),
             lambda node: str(node.target.id) + '.png'),
     }
 
@@ -110,11 +115,18 @@ def _visit_input_ann_assign(self, node, annotation):
         return None
 
     def _visit_default_dumper(self, node, dumper):
-        dump_tree = ast.parse(dumper[0].format(var_name=node.target.id))
+        if dumper[0][0] is None:
+            pre_code_body = []
+        else:
+            pre_code_body = ast.parse(dumper[0][0].format(var_name=node.target.id)).body
+        if dumper[0][1] is None:
+            post_code_body = []
+        else:
+            post_code_body = ast.parse(dumper[0][1].format(var_name=node.target.id)).body
         self.extracted_variables.append(_VariableNameTypePair(
             node.target.id, None, None, None, False, True, dumper[1](node))
         )
-        return [self.conv_AnnAssign_to_Assign(node), *dump_tree.body]
+        return [*pre_code_body, self.conv_AnnAssign_to_Assign(node), *post_code_body]
 
     def _visit_user_defined_dumper(self, node):
         load_ctx = ast.Load()

examples/new_data.png

@@ -162,7 +162,7 @@ class CWLDumpableBinaryFile(CWLDumpable):
 
 class CWLPNGPlot(CWLDumpable):
     """Use that annotation to define that after the assigment of that variable the plt.savefig() should
-    be called
+    be called.
 
     >>> import matplotlib.pyplot as plt
     >>> data = [1,2,3]
@@ -175,5 +175,33 @@ class CWLPNGPlot(CWLDumpable):
     >>> new_data: 'CWLPNGPlot' = plt.plot(data)
     >>> plt.savefig('new_data.png')
 
+
+    Note that by default if you have multiple plot statements in the same notebook will be written
+    in the same file. If you want to write them in separates you have to do it in separate figures.
+    To do that in your notebook you have to create a new figure before the plot command or use the CWLPNGFigure.
+
+    >>> import matplotlib.pyplot as plt
+    >>> data = [1,2,3]
+    >>> plt.figure()
+    >>> new_data: 'CWLPNGPlot' = plt.plot(data)
     """
     pass
+
+
+class CWLPNGFigure(CWLDumpable):
+    """The same with CWLPNGPlot but creates new figures before plotting. Use that annotation of you don't want
+    to write multiple graphs in the same image
+
+    >>> import matplotlib.pyplot as plt
+    >>> data = [1,2,3]
+    >>> new_data: 'CWLPNGPlot' = plt.plot(data)
+
+    the converter will tranform these lines to
+
+    >>> import matplotlib.pyplot as plt
+    >>> data = [1,2,3]
+    >>> plt.figure()
+    >>> new_data: 'CWLPNGPlot' = plt.plot(data)
+    >>> plt.savefig('new_data.png')
+
+    """

ipython2cwl/cwltoolextractor.py

@@ -516,3 +516,45 @@ def test_AnnotatedIPython2CWLToolConverter_CWLPNGPlot(self):
             },
             tool
         )
+
+    def test_AnnotatedIPython2CWLToolConverter_CWLPNGFigure(self):
+        code = os.linesep.join([
+            "import matplotlib.pyplot as plt",
+            "new_data: 'CWLPNGFigure' = plt.plot([1,2,3,4])",
+        ])
+        converter = AnnotatedIPython2CWLToolConverter(code)
+        new_script = converter._wrap_script_to_method(
+            converter._tree,
+            converter._variables
+        )
+        try:
+            os.remove('new_data.png')
+        except FileNotFoundError:
+            pass
+        exec(new_script)
+        locals()['main']()
+        self.assertTrue(os.path.isfile('new_data.png'))
+        os.remove('new_data.png')
+
+        tool = converter.cwl_command_line_tool()
+        self.assertDictEqual(
+            {
+                'cwlVersion': "v1.1",
+                'class': 'CommandLineTool',
+                'baseCommand': 'notebookTool',
+                'hints': {
+                    'DockerRequirement': {'dockerImageId': 'jn2cwl:latest'}
+                },
+                'arguments': ['--'],
+                'inputs': {},
+                'outputs': {
+                    'new_data': {
+                        'type': 'File',
+                        'outputBinding': {
+                            'glob': 'new_data.png'
+                        }
+                    }
+                },
+            },
+            tool
+        )