Merge pull request #918 from ReactionMechanismGenerator/draw_svg

Improve flexibility of Molecule and Group drawing algorithms

rmgpy/molecule/draw.py

@@ -58,12 +58,15 @@
 
 ################################################################################
 
-def createNewSurface(format, path=None, width=1024, height=768):
+def createNewSurface(format, target=None, width=1024, height=768):
     """
-    Create a new surface of the specified `type`: "png" for
-    :class:`ImageSurface`, "svg" for :class:`SVGSurface`, "pdf" for
-    :class:`PDFSurface`, or "ps" for :class:`PSSurface`. If the surface is to
-    be saved to a file, use the `path` parameter to give the path to the file.
+    Create a new surface of the specified `format`:
+        "png" for :class:`ImageSurface`
+        "svg" for :class:`SVGSurface`
+        "pdf" for :class:`PDFSurface`
+        "ps" for :class:`PSSurface`
+    The surface will be written to the `target` parameter , which can be a
+    path to save the surface to, or file-like object with a `write()` method.
     You can also optionally specify the `width` and `height` of the generated
     surface if you know what it is; otherwise a default size of 1024 by 768 is
     used.
@@ -76,11 +79,11 @@ def createNewSurface(format, path=None, width=1024, height=768):
     if format == 'png':
         surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, int(width), int(height))
     elif format == 'svg':
-        surface = cairo.SVGSurface(path, width, height)
+        surface = cairo.SVGSurface(target, width, height)
     elif format == 'pdf':
-        surface = cairo.PDFSurface(path, width, height)
+        surface = cairo.PDFSurface(target, width, height)
     elif format == 'ps':
-        surface = cairo.PSSurface(path, width, height)
+        surface = cairo.PSSurface(target, width, height)
     else:
         raise ValueError('Invalid value "{0}" for type parameter; valid values are "png", "svg", "pdf", and "ps".'.format(type))
     return surface
@@ -125,7 +128,7 @@ def clear(self):
         self.surface = None
         self.cr = None
 
-    def draw(self, molecule, format, path=None):
+    def draw(self, molecule, format, target=None):
         """
         Draw the given `molecule` using the given image `format` - pdf, svg, ps, or
         png. If `path` is given, the drawing is saved to that location on disk. The
@@ -225,7 +228,7 @@ def draw(self, molecule, format, path=None):
 
         # Create a dummy surface to draw to, since we don't know the bounding rect
         # We will copy this to another surface with the correct bounding rect
-        surface0 = createNewSurface(format=format, path=None)
+        surface0 = createNewSurface(format=format, target=None)
         cr0 = cairo.Context(surface0)
 
         # Render using Cairo
@@ -236,7 +239,7 @@ def draw(self, molecule, format, path=None):
         yoff = self.top
         width = self.right - self.left
         height = self.bottom - self.top
-        self.surface = createNewSurface(format=format, path=path, width=width, height=height)
+        self.surface = createNewSurface(format=format, target=target, width=width, height=height)
         self.cr = cairo.Context(self.surface)
 
         # Draw white background
@@ -245,12 +248,15 @@ def draw(self, molecule, format, path=None):
 
         self.render(self.cr, offset=(-xoff,-yoff))
 
-        if path is not None:
+        if target is not None:
             # Finish Cairo drawing
             # Save PNG of drawing if appropriate
-            ext = os.path.splitext(path)[1].lower()
-            if ext == '.png':
-                self.surface.write_to_png(path)
+            if isinstance(target, str):
+                ext = os.path.splitext(target)[1].lower()
+                if ext == '.png':
+                    self.surface.write_to_png(target)
+                else:
+                    self.surface.finish()
             else:
                 self.surface.finish()
 

rmgpy/molecule/drawTest.py

@@ -37,7 +37,7 @@ def testDrawPNG(self):
         path = 'test_molecule.png'
         if os.path.exists(path):
             os.unlink(path)
-        surface, cr, (xoff, yoff, width, height) = self.drawer.draw(self.molecule, format='png', path=path)
+        surface, cr, (xoff, yoff, width, height) = self.drawer.draw(self.molecule, format='png', target=path)
         self.assertTrue(os.path.exists(path), "File doesn't exist")
         os.unlink(path)
         self.assertIsInstance(surface, ImageSurface)
@@ -53,7 +53,7 @@ def testDrawPDF(self):
         path = 'test_molecule.pdf'
         if os.path.exists(path):
             os.unlink(path)
-        surface, cr, (xoff, yoff, width, height) = self.drawer.draw(self.molecule, format='pdf', path=path)
+        surface, cr, (xoff, yoff, width, height) = self.drawer.draw(self.molecule, format='pdf', target=path)
         self.assertIsInstance(surface, PDFSurface)
         self.assertGreater(width, height)
         os.unlink(path)
@@ -70,7 +70,7 @@ def testDrawPolycycle(self):
         if os.path.exists(path):
             os.unlink(path)
         polycycle = Molecule(SMILES="C123CC4CC1COCC2CCC34")
-        surface, cr, (xoff, yoff, width, height) = self.drawer.draw(self.molecule, format='pdf', path=path)
+        surface, cr, (xoff, yoff, width, height) = self.drawer.draw(self.molecule, format='pdf', target=path)
         self.assertIsInstance(surface, PDFSurface)
         self.assertGreater(width, height)
         os.unlink(path)

rmgpy/molecule/group.py

@@ -792,8 +792,15 @@ def _repr_png_(self):
         """
         Return a png picture of the group, useful for ipython-qtconsole.
         """
+        return self.draw('png')
+
+    def draw(self, format):
+        """
+        Use pydot to draw a basic graph of the group.
+
+        Use format to specify the desired output format, eg. 'png', 'svg', 'ps', 'pdf', 'plain', etc.
+        """
         import pydot
-        import os
 
         graph = pydot.Dot(graph_type='graph', dpi="52")
         for index, atom in enumerate(self.atoms):
@@ -810,7 +817,7 @@ def _repr_png_(self):
                     bondType = '"' + bondType + '"'
                     graph.add_edge(pydot.Edge(src=str(index1 + 1), dst=str(index2 + 1), label=bondType, fontname="Helvetica", fontsize="16"))
 
-        img = graph.create_png(prog='neato')
+        img = graph.create(prog='neato', format=format)
         return img
 
     def __getAtoms(self): return self.vertices

rmgpy/molecule/groupTest.py

@@ -1151,6 +1151,46 @@ def test_repr_png(self):
         result = group._repr_png_()
         self.assertIsNotNone(result)
 
+    def testDrawGroup(self):
+        """Test that the draw method returns the expected pydot graph."""
+        adjlist = """
+1 *1 [C,Cd,Ct,CO,CS,Cb] u1 {2,[S,D,T,B]}
+2 *2 [C,Cd,Ct,CO,CS,Cb] u0 {1,[S,D,T,B]} {3,[S,D,T,B]}
+3 *3 [C,Cd,Ct,CO,CS,Cb] u0 {2,[S,D,T,B]} {4,[S,D,T,B]}
+4 *4 [C,Cd,Ct,CO,CS,Cb] u0 {3,[S,D,T,B]}
+        """
+        # Use of tabs in the expected string is intentional
+        expected = """
+graph G {
+	graph [dpi=52];
+	node [label="\N"];
+	1	 [fontname=Helvetica,
+		fontsize=16,
+		label="*1 C,Cd,Ct,CO,CS,Cb"];
+	2	 [fontname=Helvetica,
+		fontsize=16,
+		label="*2 C,Cd,Ct,CO,CS,Cb"];
+	1 -- 2	 [fontname=Helvetica,
+		fontsize=16,
+		label="S,D,T,B"];
+	3	 [fontname=Helvetica,
+		fontsize=16,
+		label="*3 C,Cd,Ct,CO,CS,Cb"];
+	2 -- 3	 [fontname=Helvetica,
+		fontsize=16,
+		label="S,D,T,B"];
+	4	 [fontname=Helvetica,
+		fontsize=16,
+		label="*4 C,Cd,Ct,CO,CS,Cb"];
+	3 -- 4	 [fontname=Helvetica,
+		fontsize=16,
+		label="S,D,T,B"];
+}
+        """
+        group = Group().fromAdjacencyList(adjlist)
+        result = group.draw('canon')
+        self.assertEqual(''.join(result.split()), ''.join(expected.split()))
+
 ################################################################################
 
 if __name__ == '__main__':

rmgpy/molecule/molecule.py

@@ -1225,7 +1225,7 @@ def draw(self, path):
         """
         from .draw import MoleculeDrawer
         format = os.path.splitext(path)[-1][1:].lower()
-        MoleculeDrawer().draw(self, format, path=path)
+        MoleculeDrawer().draw(self, format, target=path)
 
     def _repr_png_(self):
         """

rmgpy/pdep/draw.py

@@ -362,7 +362,7 @@ def draw(self, network, format, path=None):
         width = max([rect[2]+rect[0] for rect in wellRects]) - min([rect[0] for rect in wellRects]) + 2 * padding
 
         # Draw to the final surface
-        surface = createNewSurface(format=format, path=path, width=width, height=height)
+        surface = createNewSurface(format=format, target=path, width=width, height=height)
         cr = cairo.Context(surface)
 
         # Some global settings