#!/usr/bin/python2 # -*- coding: utf-8 -*- import sys import os import shutil import re from StringIO import StringIO import yaml def _raise_err(format, *args) : raise ValueError(format % args) def _load_yaml(yaml_file) : print('load: ' + yaml_file) with open(yaml_file, 'r') as f : return yaml.load(f.read()) def _exec(cmd) : print(cmd) return os.system(cmd) class Node : all = {} keys = [] @classmethod def init(cls) : Node.all = {} Node.keys = [] for node_type in ['person', 'company'] : for node_id in os.listdir('../data/' + node_type) : yaml_file = '../data/%s/%s/brief.yaml' % (node_type, node_id) node_id = node_id.decode('utf-8') node = Node(_load_yaml(yaml_file), node_id, node_type) if node_id in Node.all : _raise_err(u'Node id conflict: "%s"!', node_id) Node.all[node_id] = node Node.keys.append(node_id) print('Node number: %d' % len(Node.all)) def __init__(self, yaml, node_id, type) : self.id = node_id self.type = type self.name = yaml['name'] if 'other_names' in yaml : # person self.other_names = yaml['other_names'] if 'sex' in yaml : # person self.sex = yaml['sex'] if 'full_name' in yaml : # company self.full_name = yaml['full_name'] self.birth = yaml['birth'] self.death = yaml['death'] self.desc = yaml['desc'] self.links = yaml['links'] class Relation : all = {} keys = [] @classmethod def init(cls) : for family_id in os.listdir('../data/family/') : family_id = family_id.replace('.yaml', '') yaml_file = '../data/family/%s.yaml' % (family_id,) family_id = family_id.decode('utf-8') if family_id not in Node.all : _raise_err(u'Invalid family name: "%s"!', family_id) yaml = _load_yaml(yaml_file) for lst in yaml['relations'] : relation = Relation(lst) Relation.all[relation.name] = relation Relation.keys.append(relation.name) print('Relation number: %d' % len(Relation.all)) def __init__(self, lst) : self.node_from = lst[0] self.node_to = lst[1] self.desc = lst[2] self.name = self.node_from + '->' + self.node_to if self.name in Relation.all : _raise_err(u'Relation name conflict: "%s"!', self.name) if self.node_from not in Node.all : _raise_err(u'Invalid relation "from" attr: "%s"!', self.node_from) if self.node_to not in Node.all : _raise_err(u'Invalid relation "to" attr": "%s"!', self.node_to) class Family : all = {} keys = [] @classmethod def init(cls) : for family_id in os.listdir('../data/family/') : family_id = family_id.replace('.yaml', '') yaml_file = '../data/family/%s.yaml' % (family_id,) family_id = family_id.decode('utf-8') if family_id not in Node.all : _raise_err(u'Invalid family name: "%s"!', family_id) family = Family(_load_yaml(yaml_file)) Family.all[family_id] = family Family.keys.append(family_id) print('Family number: %d' % len(Family.all)) def __init__(self, yaml) : self.name = yaml['name'] self.inner = yaml['inner'] self.outer = yaml['outer'] self.members = [self.name] + self.inner + self.outer for name in self.members : if name not in Node.all : _raise_err(u'Invalid family members: "%s"!', name) class Graph : def __init__(self, yaml) : self._name = yaml['name'] self._families = yaml['families'] self._families.reverse() self._nodes = [] self._relations = [] for f in self._families : family = Family.all[f] for n in family.members : if n not in self._nodes : self._nodes.append(n) for r in Relation.keys : relation = Relation.all[r] if relation.node_from in family.members \ and relation.node_to in family.members \ and r not in self._relations : self._relations.append(r) def __unicode__(self) : output = StringIO() for n in self._nodes : output.write(self._dot_node(n)) output.write('\n') for r in self._relations : output.write(self._dot_relation(r)) output.write('\n') if len(self._families) > 1 : for f in self._families : output.write(self._dot_sub_graph(f)) template = u''' digraph %s { \trankdir = "LR"; \tranksep = 0.5; \tlabel = "%s"; \tlabelloc = "t"; \tfontsize = "24"; \tfontname = "SimHei"; \tgraph [style="filled", color="lightgrey"]; \tnode [fontname="SimSun"]; \tedge [fontname="SimSun"]; %s } ''' return template % (self._name, self._name, output.getvalue()) def _node_color(self, node) : if u'company' == node.type : return u'green' else : return (u'blue' if u'M'==node.sex else u'red') def _dot_node(self, node_id) : node = Node.all[node_id] template = u'''\t%s [shape="%s", color="%s", ''' \ u'''label=<
%s%s |
%s |
%s |