Improve tagdag command

- Support limiting depth of the graph with `--depth`
- Slight performance improvement due to caching

sok/management/commands/tagdag.py

@@ -1,5 +1,6 @@
 import html
 
+from functools import lru_cache
 from typing import Optional, Set, Tuple
 
 from django.core.management.base import BaseCommand, CommandParser
@@ -12,15 +13,28 @@ class Command(BaseCommand):
 	def echo(self, msg: str, nl: bool = True):
 		self.stdout.write(msg, ending='\n' if nl else '')
 
+	@lru_cache
+	def transitive_publications(self, tag: Tag):
+		return tag.transitive_publications
+
+	@lru_cache
+	def num_publications(self, tag: Tag):
+		return len(self.transitive_publications(tag))
+
 	def add_node(
 		self,
 		node: Tag,
 		publication: Optional[Publication] = None,
 		threshold: int = 0,
 		include_publications: bool = False,
+		max_depth: int = 0,
+		depth: int = 0,
 	):
-		publications = node.transitive_publications
-		num = len(publications)
+		if 0 < max_depth and max_depth < depth:
+			return
+
+		publications = self.transitive_publications(node)
+		num = self.num_publications(node)
 
 		if node.pk in self.nodes:
 			return  # Already printed this node
@@ -49,7 +63,8 @@ class Command(BaseCommand):
 				implicit = True
 		self.echo('",')
 		if 0 == num:
-			self.echo("\t\tcolor=red,")
+			self.echo("\t\tcolor=firebrick2,")
+			self.echo("\t\tfontcolor=firebrick3,")
 		if implicit:
 			self.echo("\t\tcolor=gainsboro,")
 			self.echo("\t\tfontcolor=gray,")
@@ -57,7 +72,7 @@ class Command(BaseCommand):
 
 		self.nodes.add(node.pk)
 		for predecessor in node.implied_by.all():
-			self.add_node(predecessor, publication, threshold, include_publications)
+			self.add_node(predecessor, publication, threshold, include_publications, max_depth, depth + 1)
 
 	def add_edge(self, node: Tag):
 		for predecessor in node.implied_by.all():
@@ -69,7 +84,10 @@ class Command(BaseCommand):
 			if edge[::-1] in self.graph:
 				self.stderr.write(self.style.ERROR(f"CYCLE: '{node}' <-> '{predecessor}'"))
 			self.graph.add(edge)
-			self.echo(f"\tT{predecessor.pk} -> T{node.pk};")
+			self.echo(f"\tT{predecessor.pk} -> T{node.pk}", nl=False)
+			if 0 == self.num_publications(predecessor):
+				self.echo(" [color=firebrick2]", nl=False)
+			self.echo(";")
 			self.add_edge(predecessor)
 
 	def graphviz(
@@ -78,6 +96,7 @@ class Command(BaseCommand):
 		publication: Optional[Publication] = None,
 		threshold: int = 0,
 		include_publications: bool = False,
+		max_depth: int = 0,
 	):
 		self.echo("digraph G {")
 		self.echo("\trankdir = RL;")
@@ -86,9 +105,9 @@ class Command(BaseCommand):
 		# Add nodes
 		if root is None:
 			for tag in Tag.objects.filter(implies__isnull=True):
-				self.add_node(tag, publication, threshold, include_publications)
+				self.add_node(tag, publication, threshold, include_publications, max_depth)
 		else:
-			self.add_node(root, publication, threshold, include_publications)
+			self.add_node(root, publication, threshold, include_publications, max_depth)
 
 		# Add edges
 		if root is None:
@@ -105,11 +124,13 @@ class Command(BaseCommand):
 		parser.add_argument('--root', default=None)
 		parser.add_argument('--include-publications', action='store_true')
 		parser.add_argument('--threshold', type=int, default=0)
+		parser.add_argument('--depth', type=int, default=0)
 		parser.add_argument('publication', nargs='?')
 
 	def handle(self, *args, **options) -> None:
 		include_publications: bool = options['include_publications']
 		threshold: int = options['threshold']
+		max_depth: int = options['depth']
 
 		root: Optional[Tag] = None
 		if tag_name := options.get('root', None):
@@ -121,4 +142,4 @@ class Command(BaseCommand):
 
 		self.graph: Set[Tuple[int, int]] = set()
 		self.nodes: Set[int] = set()
-		self.graphviz(root, publication, threshold, include_publications)
+		self.graphviz(root, publication, threshold, include_publications, max_depth)