rdkit.Chem.Draw.aggCanvas

Package rdkit :: Package Chem :: Package Draw :: Module aggCanvas |
[hide private]
[frames] | no frames]

# Source Code for Module rdkit.Chem.Draw.aggCanvas

      1  #
      2  #  Copyright (C) 2008 Greg Landrum
      3  #
      4  #   @@ All Rights Reserved @@
      5  #  This file is part of the RDKit.
      6  #  The contents are covered by the terms of the BSD license
      7  #  which is included in the file license.txt, found at the root
      8  #  of the RDKit source tree.
      9  #
     10  import os
     11  import re
     12
     13  from aggdraw import Brush, Pen
     14  from aggdraw import Draw
     15  from aggdraw import Font
     16  from rdkit import RDConfig
     17  from rdkit.Chem.Draw.canvasbase import CanvasBase
     18
     19  faceMap = {'sans': os.path.join(RDConfig.RDCodeDir, 'Chem', 'Draw', 'FreeSans.ttf')}
     20
     21
     22 -def convertColor(color):
     23    color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))
     24    return color
     25
     26
     27 -class Canvas(CanvasBase):
     28    # fonts appear smaller in aggdraw than with cairo
     29    # fix that here:
     30    fontScale = 1.2
     31
     32 -  def __init__(self,
     33                 img=None,
     34                 imageType=None,  # determines file type
     35                 fileName=None,  # if set determines output file name
     36                 size=None, ):
     37      if img is None:
     38        try:
     39          import Image
     40        except ImportError:
     41          from PIL import Image
     42        if size is None:
     43          raise ValueError('please provide either an image or a size')
     44        img = Image.new('RGBA', size, "white")
     45      self.image = img
     46      self.draw = Draw(img)
     47      self.draw.setantialias(True)
     48      if size is None:
     49        self.size = self.draw.size
     50      else:
     51        self.size = size
     52      if imageType and imageType not in ('png', 'jpg'):
     53        raise ValueError('unsupported image type for agg canvas')
     54      self.drawType = imageType
     55      self.fileName = fileName
     56
     57 -  def _doLine(self, p1, p2, pen, **kwargs):
     58      if kwargs.get('dash', (0, 0)) == (0, 0):
     59        self.draw.line((p1[0], p1[1], p2[0], p2[1]), pen)
     60      else:
     61        dash = kwargs['dash']
     62        pts = self._getLinePoints(p1, p2, dash)
     63
     64        currDash = 0
     65        dashOn = True
     66        while currDash < (len(pts) - 1):
     67          if dashOn:
     68            p1 = pts[currDash]
     69            p2 = pts[currDash + 1]
     70            self.draw.line((p1[0], p1[1], p2[0], p2[1]), pen)
     71          currDash += 1
     72          dashOn = not dashOn
     73
     74 -  def addCanvasLine(self, p1, p2, color=(0, 0, 0), color2=None, **kwargs):
     75      if color2 and color2 != color:
     76        mp = (p1[0] + p2[0]) / 2., (p1[1] + p2[1]) / 2.
     77        color = convertColor(color)
     78        self._doLine(p1, mp, Pen(color, kwargs.get('linewidth', 1)), **kwargs)
     79        color2 = convertColor(color2)
     80        self._doLine(mp, p2, Pen(color2, kwargs.get('linewidth', 1)), **kwargs)
     81      else:
     82        color = convertColor(color)
     83        self._doLine(p1, p2, Pen(color, kwargs.get('linewidth', 1)), **kwargs)
     84
     85 -  def addCanvasText(self, text, pos, font, color=(0, 0, 0), **kwargs):
     86      orientation = kwargs.get('orientation', 'E')
     87      color = convertColor(color)
     88      aggFont = Font(color, faceMap[font.face], size=font.size * self.fontScale)
     89
     90      blocks = list(re.finditer(r'\<(.+?)\>(.+?)\</\1\>', text))
     91      w, h = 0, 0
     92      supH = 0
     93      subH = 0
     94      if not len(blocks):
     95        w, h = self.draw.textsize(text, aggFont)
     96        tw, th = w, h
     97        offset = w * pos[2]
     98        dPos = pos[0] - w / 2. + offset, pos[1] - h / 2.
     99        self.draw.text(dPos, text, aggFont)
    100      else:
    101        dblocks = []
    102        idx = 0
    103        for block in blocks:
    104          blockStart, blockEnd = block.span(0)
    105          if blockStart != idx:
    106            # untagged text:
    107            tblock = text[idx:blockStart]
    108            tw, th = self.draw.textsize(tblock, aggFont)
    109            w += tw
    110            h = max(h, th)
    111            dblocks.append((tblock, '', tw, th))
    112          fmt = block.groups()[0]
    113          tblock = block.groups()[1]
    114          if fmt in ('sub', 'sup'):
    115            lFont = Font(color, faceMap[font.face], size=0.8 * font.size * self.fontScale)
    116          else:
    117            lFont = aggFont
    118          tw, th = self.draw.textsize(tblock, lFont)
    119          w += tw
    120          if fmt == 'sub':
    121            subH = max(subH, th)
    122          elif fmt == 'sup':
    123            supH = max(supH, th)
    124          else:
    125            h = max(h, th)
    126          dblocks.append((tblock, fmt, tw, th))
    127          idx = blockEnd
    128        if idx != len(text):
    129          # untagged text:
    130          tblock = text[idx:]
    131          tw, th = self.draw.textsize(tblock, aggFont)
    132          w += tw
    133          h = max(h, th)
    134          dblocks.append((tblock, '', tw, th))
    135
    136        supH *= 0.5
    137        subH *= 0.5
    138        h += supH + subH
    139        offset = w * pos[2]
    140        dPos = [pos[0] - w / 2. + offset, pos[1] - h / 2.]
    141        if orientation == 'W':
    142          dPos = [pos[0] - w + offset, pos[1] - h / 2.]
    143        elif orientation == 'E':
    144          dPos = [pos[0] + offset, pos[1] - h / 2.]
    145        else:
    146          dPos = [pos[0] - w / 2 + offset, pos[1] - h / 2.]
    147
    148        if supH:
    149          dPos[1] += supH
    150        for txt, fmt, tw, th in dblocks:
    151          tPos = dPos[:]
    152          if fmt == 'sub':
    153            tPos[1] += subH
    154          elif fmt == 'sup':
    155            tPos[1] -= supH
    156          if fmt in ('sub', 'sup'):
    157            lFont = Font(color, faceMap[font.face], size=0.8 * font.size * self.fontScale)
    158          else:
    159            lFont = aggFont
    160          self.draw.text(tPos, txt, lFont)
    161          dPos[0] += tw
    162      return (tw + th * .4, th + th * .4, offset)
    163
    164 -  def addCanvasPolygon(self, ps, color=(0, 0, 0), fill=True, stroke=False, **kwargs):
    165      if not fill and not stroke:
    166        return
    167      dps = []
    168      for p in ps:
    169        dps.extend(p)
    170      color = convertColor(color)
    171      brush = None
    172      pen = None
    173      if fill:
    174        brush = Brush(color)
    175      if stroke:
    176        pen = Pen(color)
    177      self.draw.polygon(dps, pen, brush)
    178
    179 -  def addCanvasDashedWedge(self, p1, p2, p3, dash=(2, 2), color=(0, 0, 0), color2=None, **kwargs):
    180      pen = Pen(color, kwargs.get('linewidth', 1))
    181      dash = (3, 3)
    182      pts1 = self._getLinePoints(p1, p2, dash)
    183      pts2 = self._getLinePoints(p1, p3, dash)
    184
    185      if len(pts2) < len(pts1):
    186        pts2, pts1 = pts1, pts2
    187
    188      for i in range(len(pts1)):
    189        self.draw.line((pts1[i][0], pts1[i][1], pts2[i][0], pts2[i][1]), pen)
    190
    191 -  def flush(self):
    192      self.draw.flush()
    193      if self.fileName:
    194        self.image.save(self.fileName)
    195