https://bugs.launchpad.net/zim/+bug/1472408/+attachment/4427202/+files/treemapeditor.py

treemapeditor.py

# -*- coding: utf-8 -*-
#
# treemapeditor.py
#
# This is a plugin for Zim, which creates treeMap graph based on tabbed keywords.
# Initial version comes from the following source: http://blog.ynema.com/?p=192
#
#
# Author: NorfCran <norfcran@gmail.com>
# Date: 2015-07-09
# Copyright (c) 2015, released under the GNU GPL v2 or higher
#
#

import logging
try:
  import pydot
  has_dotmod = True
except ImportError:
  has_dotmod = False
     

from zim.plugins.base.imagegenerator import ImageGeneratorPlugin, ImageGeneratorClass
from zim.fs import File, TmpFile
from zim.config import data_file
from zim.applications import Application, ApplicationError

# TODO put these commands in preferences
dotcmd = ('twopi', '-Tpng', '-o')

logger = logging.getLogger('zim.plugins')

class InsertTreeMapPlugin(ImageGeneratorPlugin):

  plugin_info = {
    'name': _('Insert TreeMap'), # T: plugin name
    'description': _('''\
This plugin provides a hierachical editor for zim based on Twopi (part of Graphviz).
'''), # T: plugin description
        'help': 'Plugins:TreeMap Editor',
    'author': 'NorfCran',
  }

  object_type = 'treemap'
  short_label = _('TreeMap') # T: menu item
  insert_label = _('Insert TreeMap') # T: menu item
  edit_label = _('_Edit TreeMap') # T: menu item
  syntax = None

  #@classmethod
  #def check_dependencies(klass):
  #  has_dotmod = True
  #  try:
  #      import pydot
  #  except ImportError:
  #      has_dotmod = True
  #  return has_dotmod, [("GraphViz", has_dotmod, True)]

  @classmethod
  def check_dependencies(klass):
     has_dotcmd = Application(dotcmd).tryexec()
     return (has_dotmod and has_dotcmd), \
        [("pydot", has_dotmod, True), ("GraphViz", has_dotcmd, True)]


class TreeMapGenerator(ImageGeneratorClass):

  uses_log_file = False

  object_type = 'treemap'
  scriptname = 'treemap.tab'
  imagename = 'treemap.png'

  def __init__(self, plugin):
    ImageGeneratorClass.__init__(self, plugin)
    self.dotfile = TmpFile(self.scriptname)
    self.dotfile.touch()
    self.pngfile = File(self.dotfile.path[:-4] + '.png') # len('.dot') == 4

  def generate_image(self, text):
    # string has to be treated differently, if list is passed do nothing
    if type(text) is str:
      text = text.split("\n")
    if len(text) == 1:
      text = [text[0],"\tnewline + TAB + keyword"]
    self.dotfile.write(self.createGraphFromEdges(self.createEdges(text)))
    # Call GraphViz
    try:
      dot = Application(dotcmd)
      dot.run((self.pngfile, self.dotfile))
    except ApplicationError:
      return None, None # Sorry, no log
    else:
      if self.pngfile.exists():
        return self.pngfile, None
      else:
        return None, None

  def cleanup(self):
    self.dotfile.remove()
    self.pngfile.remove()

  # Configuration variables for the final treemap
  ranksep = '2 equally' # distance between nodes (suggested values: 1-3) => (graphviz varialbe)
  overlap='true' # whether the nodes overlap each other (true or false) => (graphviz varialbe)
  splines='true' # whether the lines are straight or curved (true or false) => (graphviz varialbe)
  gen_dot = True # whether the dot file should be generated => (python varialbe)
  plain_txt = True

  # the list can be managed by a manually defined array of colors
  # http://www.graphviz.org/doc/info/colors.html
  def generateColors(self):
    colors_nodes = self.initArrayOfcolors('black', ['grey', '#fdd49e', '#fc8d59', '#d7301f', '#b00300', '#700f00'])
    colors_font = self.initArrayOfcolors('white', ['black', 'black', 'black'])
    return (colors_nodes, colors_font)

  def initArrayOfcolors(self, default, colors):
    # generates an array of 12 colors
    colors_nodes = [default for x in range(12)]
    return colors + colors_nodes[len(colors):]

  def getAttributes(self, node, level):
    # provides access to colors
    colors_nodes, colors_font = self.generateColors()
    # investigates edge
    # root node
    if level is 0:
      return {'shape':'box, filled',
      'fillcolor':'gray',
      'style':'"filled"',
      'penwidth':'3',
      'color':colors_nodes[level + 1],
      'fontcolor':'black',
      'fontsize':'32'}
    # default edge
    if type(node) is tuple:
      return {
      'penwidth':'2',
      'color':colors_nodes[level + 1]}
    # investigates node
    # remove whitespaces
    node = node.strip()
    # __underlined__
    if node.startswith("\"__") and node.endswith("__\""):
      return {'shape':'box',
      'style':'"rounded,filled"',
      'penwidth':'2',
      'label':node[2:-2],
      'fillcolor':'yellow',
      'color':colors_nodes[level],
      'fontcolor':'black'}
    # **bold**
    if node.startswith("\"**") and node.endswith("**\""):
      return {'shape':'box',
      'style':'"rounded, filled"',
      'penwidth':'2',
      'label':node[2:-2],
      'color':'black',
      'fillcolor':'red',
      'fontcolor':'black',
      'fontsize':'15'}
    # //italic//
    if node.startswith("\"//") and node.endswith("//\""):
      return {'shape':'box',
      'style':'"rounded, filled"',
      'penwidth':'2',
      'label':node[2:-2],
      'color':'black',
      'fillcolor':'green',
      'fontcolor':'black', }
    # standard node
    else:
      return {'shape':'box',
      'style':'filled',
      'penwidth':'2',
      'fillcolor':colors_nodes[level],
      'color':colors_nodes[level + 1],
      'fontcolor':colors_font[level]}

  def createEdges(self, lines):
    edge_list = ['' for x in range(50)]
    edges = []
    for line in lines:
      # identify right indent (tabs excluding tabs in the text)
      tabs_all = line.count('\t')
      line = line.strip().replace(":", "")
      pos = tabs_all - line.count('\t')
      # removes comments (tabbed text)
      index = line.find("\t")
      if index is not -1: line = line[:index]
      # assign extracted text from a line
      edge_list[pos] = "\"%s\"" % line
      if pos:
        edges.append((edge_list[pos - 1], edge_list[pos], pos - 1))
    return edges

  def createGraphFromEdges(self, edges):
    # defines graph (according to configuration variables)
    g = pydot.Dot(splines=self.splines, overlap=self.overlap, ranksep=self.ranksep, root="%s" % edges[0][0])
    # iterates through edges
    for edge in edges:
      g.add_node(pydot.Node(edge[0], **self.getAttributes(edge[0], edge[2])))
      g.add_node(pydot.Node(edge[1], **self.getAttributes(edge[1], edge[2] + 1)))
      g.add_edge(pydot.Edge(edge[0], edge[1], **self.getAttributes((edge[0], edge[0]), edge[2])))
    # in a case that there is only one mindmap on a page there is no reason for extension of a mm_title to a mm_title_rootnode
    return g.create_dot()