class Command(BaseCommand): def handle(self, *apps, **options): self.graph = MigrationLoader(None).graph comment = options['comment'] self.picture = Digraph(comment=comment) save_loc = options['filename'] self.render(save_loc) def add_arguments(self, parser): parser.add_argument('--comment', help='optional comments/captions for the picture') parser.add_argument('filename', nargs='?', help='a filename to write GraphViz contents to') def _create_digraph(self): nodes = sorted(self.graph.nodes.values(), key=self._get_tuple) for node in nodes: self._add_node(node) for node in nodes: self._add_dependencies(node) @staticmethod def _style_label(tupled_node): return '/'.join(tupled_node) @staticmethod def _get_tuple(node): return (node.app_label, node.name) def _add_node(self, node): node_label = self._style_label(self._get_tuple(node)) self.picture.node(node_label, node_label) def _add_edges(self, nodeTo, nodeFrom): self.picture.edge(self._style_label(nodeFrom), self._style_label(nodeTo)) def _add_dependencies(self, node): for dep in node.dependencies: if dep[1] == '__first__': self._add_edges( self.graph.root_nodes(dep[0])[0], self._get_tuple(node)) elif dep[1] == '__latest__': self._add_edges( self.graph.leaf_nodes(dep[0])[0], self._get_tuple(node)) else: self._add_edges(dep, self._get_tuple(node)) def render(self, save_loc): self._create_digraph() if save_loc: self.picture.render(save_loc, view=False) else: with NamedTemporaryFile() as temp: self.picture.render(temp.name, view=True)
class Command(BaseCommand): def handle(self, *apps, **options): self._censor_cache = {} self._censor_enabled = bool(options['censor']) if self._censor_enabled: random.seed(options['random_seed']) self.graph = MigrationLoader(None).graph comment = options['comment'] self.picture = Digraph(comment=comment) save_loc = options['filename'] self._render(save_loc) def add_arguments(self, parser): parser.add_argument('--comment', help='optional comments/captions for the picture') parser.add_argument('--censor', action='store_true', help='censor node names (e.g. for publishing)') parser.add_argument('--random-seed', default=0, type=int, help='random seed (default: %(default)s)') parser.add_argument('filename', nargs='?', help='a filename to write GraphViz contents to') def _create_digraph(self): nodes = sorted(self.graph.nodes.values(), key=self._get_tuple) for node in nodes: self._add_node(node) for node in nodes: self._add_dependencies(node) @staticmethod def _censor(text): res = [] for word in text.split('_'): if word not in ('auto', 'initial', 'squashed'): chars = [] for c in word: if c not in '0123456789_': c = chr(random.randint(ord('a'), ord('z'))) chars.append(c) word = ''.join(chars) res.append(word) return '_'.join(res) def _censor_using_cache(self, text): try: return self._censor_cache[text] except KeyError: while True: censored = self._censor(text) if censored not in self._censor_cache.values(): break self._censor_cache[text] = censored return censored def _style_label(self, tupled_node): if self._censor_enabled: tupled_node = [self._censor_using_cache(e) for e in tupled_node] return '/'.join(tupled_node) @staticmethod def _get_tuple(node): return (node.app_label, node.name) def _add_node(self, node): node_label = self._style_label(self._get_tuple(node)) self.picture.node(node_label, node_label) def _add_edges(self, nodeTo, nodeFrom): self.picture.edge(self._style_label(nodeFrom), self._style_label(nodeTo)) def _add_dependencies(self, node): for dep in node.dependencies: if dep[1] == '__first__': self._add_edges( self.graph.root_nodes(dep[0])[0], self._get_tuple(node)) elif dep[1] == '__latest__': self._add_edges( self.graph.leaf_nodes(dep[0])[0], self._get_tuple(node)) else: self._add_edges(dep, self._get_tuple(node)) def _render(self, save_loc): self._create_digraph() if save_loc: self.picture.render(save_loc, view=False) else: with NamedTemporaryFile() as temp: self.picture.render(temp.name, view=True)