Fix and improve tagdag command

- Allows exporting a tag DAG for a specific publication
- Allows exporting the complete tag DAG. Before a root node had to be
  specified.
- Allows limiting tag DAG to specified number of relevant publications.
- Use record nodes.

README.md

@@ -85,7 +85,10 @@ Tags and citations can be exported in DOT format, which can be rendered using Gr
 ./manage.py citations --min-citations 10 > citations.dot && dot -Tsvg -ocitations.svg citations.dot
 
 # Render tag DAG
-./manage.py tagtag --root 'Literature on Systematic Mapping' > sysmap.dot && dot -Tsvg -osysmap.svg sysmap.dot
+./manage.py tagtag --root 'Literature on Systematic Mapping' --threshold 1 > sysmap.dot && dot -Tsvg -osysmap.svg sysmap.dot
+
+# Render classification for a single publication
+./manage.py tagtag 'DBLP:conf/ease/PetersenFMM08' > petersen08.dot && dot -Tsvg -opetersen08.svg petersen08.dot
 
 # Render with TIKZ
 pip install dot2tex

sok/management/commands/tagdag.py

@@ -1,8 +1,10 @@
-from typing import Set, Tuple
+import html
+
+from typing import Optional, Set, Tuple
 
 from django.core.management.base import BaseCommand, CommandParser
 
-from sok.models import Tag
+from sok.models import Publication, Tag
 
 
 class Command(BaseCommand):
@@ -10,29 +12,98 @@ class Command(BaseCommand):
 	def echo(self, msg: str):
 		self.stdout.write(msg)
 
-	# BaseCommand
+	def add_node(
+		self,
+		node: Tag,
+		publication: Optional[Publication] = None,
+		threshold: int = 0,
+		include_publications: bool = False,
+	):
+		publications = node.transitive_publications
+		num = len(publications)
 
-	def add_arguments(self, parser: CommandParser):
-		parser.add_argument('--root', default='CAPI Misuse')
+		if node.pk in self.nodes:
+			return  # Already printed this node
+		if num < threshold:
+			return
+		if not (publication is None or publication in publications):
+			return
 
-	def _graphviz(self, root: Tag) -> None:
-		for tag in root.implied_by.all():
-			edge = (tag.pk, root.pk)
+		name = html.escape(node.name)
+		self.echo(f"\tT{node.pk} [")
+		if include_publications:
+			pubs = ','.join([str(t.pk) for t in publications])
+			self.echo(f'\t\tlabel="{name}|{{{num}|{pubs}}}",')
+		else:
+			self.echo(f'\t\tlabel="{name}|{num}",')
+		if 0 == num:
+			self.echo("\t\tcolor=red,")
+		self.echo("\t];")
+
+		self.nodes.add(node.pk)
+		for predecessor in node.implied_by.all():
+			self.add_node(predecessor, publication, threshold, include_publications)
+
+	def add_edge(self, node: Tag):
+		for predecessor in node.implied_by.all():
+			if predecessor.pk not in self.nodes:
+				continue
+			edge = (predecessor.pk, node.pk)
 			if edge in self.graph:
 				continue
 			if edge[::-1] in self.graph:
-				self.stderr.write(self.style.ERROR(f"CYCLE: '{root}' <-> '{tag}'"))
+				self.stderr.write(self.style.ERROR(f"CYCLE: '{node}' <-> '{predecessor}'"))
 			self.graph.add(edge)
-			self._graphviz(tag)
-			self.echo(f'\t"{tag}" -> "{root}";')
+			self.echo(f'\t"T{predecessor.pk}" -> "T{node.pk}";')
+			self.add_edge(predecessor)
 
-	def graphviz(self, root: Tag) -> None:
+	def graphviz(
+		self,
+		root: Optional[Tag] = None,
+		publication: Optional[Publication] = None,
+		threshold: int = 0,
+		include_publications: bool = False,
+	):
 		self.echo("digraph G {")
 		self.echo("\trankdir = RL;")
-		self._graphviz(root)
+		self.echo("\tnode [shape=record];")
+
+		# Add nodes
+		if root is None:
+			for tag in Tag.objects.filter(implies__isnull=True):
+				self.add_node(tag, publication, threshold, include_publications)
+		else:
+			self.add_node(root, publication, threshold, include_publications)
+
+		# Add edges
+		if root is None:
+			for tag in Tag.objects.filter(implies__isnull=True):
+				self.add_edge(tag)
+		else:
+			self.add_edge(root)
+
 		self.echo("}")
 
+	# BaseCommand
+
+	def add_arguments(self, parser: CommandParser):
+		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('publication', nargs='?')
+
 	def handle(self, *args, **options) -> None:
+		include_publications: bool = options['include_publications']
+		threshold: int = options['threshold']
+
+		root: Optional[Tag] = None
+		if tag_name := options.get('root', None):
+			root = Tag.objects.get(name=tag_name)
+
+		publication: Optional[Publication] = None
+		if cite_key := options.get('publication', None):
+			publication = Publication.objects.get(cite_key=cite_key)
+
 		self.graph: Set[Tuple[int, int]] = set()
-		root = Tag.objects.get(name=options['root'])
-		self.graphviz(root)
+		self.nodes: Set[int] = set()
+		self.graphviz(root, publication, threshold, include_publications)