Merge pull request #15 from giannisdoukas/build-in-plot

Add build in support for matplotlib figures

examples/requirements.txt

@@ -0,0 +1,2 @@
+pandas
+matplotlib

ipython2cwl/cwltoolextractor.py

@@ -15,7 +15,7 @@
 from nbformat.notebooknode import NotebookNode  # type: ignore
 
 from .iotypes import CWLFilePathInput, CWLBooleanInput, CWLIntInput, CWLStringInput, CWLFilePathOutput, \
-    CWLDumpableFile, CWLDumpableBinaryFile, CWLDumpable
+    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:
@@ -64,9 +64,21 @@ class AnnotatedVariablesExtractor(ast.NodeTransformer):
     }
 
     dumpable_mapper = {
-        (CWLDumpableFile.__name__,): "with open('{var_name}', 'w') as f:\n\tf.write({var_name})",
-        (CWLDumpableBinaryFile.__name__,): "with open('{var_name}', 'wb') as f:\n\tf.write({var_name})",
+        (CWLDumpableFile.__name__,): (
+            (None, "with open('{var_name}', 'w') as f:\n\tf.write({var_name})",),
+            lambda node: node.target.id
+        ),
+        (CWLDumpableBinaryFile.__name__,): (
+            (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__,): (
+            (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'),
     }
 
     def __init__(self, *args, **kwargs):
@@ -110,12 +122,18 @@ def _visit_input_ann_assign(self, node, annotation):
         return None
 
     def _visit_default_dumper(self, node, dumper):
-        dump_tree = ast.parse(dumper.format(var_name=node.target.id))
-        self.to_dump.append(dump_tree.body)
+        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, node.target.id)
+            node.target.id, None, None, None, False, True, dumper[1](node))
         )
-        return self.conv_AnnAssign_to_Assign(node)
+        return [*pre_code_body, self.conv_AnnAssign_to_Assign(node), *post_code_body]
 
     def _visit_user_defined_dumper(self, node):
         load_ctx = ast.Load()

ipython2cwl/iotypes.py

@@ -158,3 +158,50 @@ class CWLDumpableBinaryFile(CWLDumpable):
     and at the CWL, the data, will be mapped as a output.
     """
     pass
+
+
+class CWLPNGPlot(CWLDumpable):
+    """Use that annotation to define that after the assigment of that variable the plt.savefig() should
+    be called.
+
+    >>> 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]
+    >>> 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 :class:`~ipython2cwl.iotypes.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')
+
+    """

test-requirements.txt

@@ -4,6 +4,7 @@ coveralls>=2.0.0
 virtualenv>=3.1.0
 gitpython>=3.1.3
 docker>=4.2.1
-git+https://github.com/giannisdoukas/cwltool.git#egg=cwltool
+cwltool==3.0.20200706173533
 pandas==1.0.5
 mypy
+matplotlib

tests/test_cwltoolextractor.py

@@ -474,3 +474,87 @@ def test_AnnotatedIPython2CWLToolConverter_custom_dumpables(self):
                 os.remove(f)
             except FileNotFoundError:
                 pass
+
+    def test_AnnotatedIPython2CWLToolConverter_CWLPNGPlot(self):
+        code = os.linesep.join([
+            "import matplotlib.pyplot as plt",
+            "new_data: 'CWLPNGPlot' = 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
+        )
+
+    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
+        )

tests/test_system_tests.py

@@ -29,8 +29,8 @@ def test_repo2cwl(self):
         self.assertListEqual(['example1.cwl'], [f for f in os.listdir(output_dir) if not f.startswith('.')])
 
         with open(os.path.join(output_dir, 'example1.cwl')) as f:
-            print(20 * '=')
             print('workflow file')
+            print(20 * '=')
             print(f.read())
             print(20 * '=')