Glottolog language tree visualizer

langtree.py

import re
import enum
import pathlib
import tempfile
import dataclasses
import argparse
import webbrowser
import typing

import Bio.Phylo
import dendroviz

class LangLevel(enum.StrEnum):
	LANGUAGE = "l"
	FAMILY = "f"
	DIALECT = "d"

@dataclasses.dataclass
class GlottoLang:
	"""
	Represents a parsed Glottolog node label.
	
	name - the common name of the language/family.
	glottocode - the unique and stable Glottolog identifier.
	iso_code - the ISO 639-3 code; optional.
	level - the classification level (language, family, or dialect); optional.
	"""
	name: str
	glottocode: str
	iso_code: str | None = None
	level: LangLevel | None = None
	
	_pattern = re.compile(r"""
		^							# Start of line
		(?P<name>.+?)				# Name: non-greedy match
		\s							# Mandatory space
		\[(?P<glottocode>.*?)\]		# Glottocode in brackets
		(?:\[(?P<iso_code>.*?)\])?	# Optional ISO code in brackets
		(?:-(?P<level>.)-)?			# Optional level flag
		$							# End of line
	""", re.VERBOSE)
	
	@classmethod
	def from_string(cls, line: str) -> "GlottoLang":
		if not isinstance(line, str):
			raise TypeError(f"Expected string for parsing, got {type(line).__name__}")
		
		if match := cls._pattern.match(line):
			data = match.groupdict()
			if data["level"]:
				data["level"] = LangLevel(data["level"])
			return cls(**data)
		
		raise ValueError(f"String does not match Glottolog syntax: '{line}'")


def draw_language_tree(
	tree_file_path: pathlib.Path | str,
	name: str | None = None,
	glottocode: str | None = None,
	**tree_gen_kwargs
):
	if not name and not glottocode:
		raise ValueError("You must provide a name or a glottocode.")
	
	trees = list(Bio.Phylo.parse(tree_file_path, "newick"))
	clade = find_clade_in_forest(trees, name=name, glottocode=glottocode)
	
	name, glottocode = _normalize_search_params(clade)
	
	_clean_node_names(clade)
	
	output_path = pathlib.Path(f"{name} [{glottocode}]").with_suffix('.svg')
	save_tree(clade, output_path, **tree_gen_kwargs)


def save_tree(clade: Bio.Phylo.BaseTree.Clade, output_path: pathlib.Path | str, **tree_gen_kwargs):
	with tempfile.TemporaryDirectory() as tmp_dir:
		tmp_path = pathlib.Path(tmp_dir) / "tree.nwk"
		Bio.Phylo.write(clade, tmp_path, "newick")
	
		generator = dendroviz.DendrogramGenerator()
		generator.generate_tree(
			tmp_path,
			input_format="newick",
			output_svg=output_path,
			show_labels=True,
			options=dendroviz.LayoutOptions(
				colour_mode="palette",
				palette="set2",
				palette_depth=2,
			),
			**tree_gen_kwargs
		)
		
	print(f"Visualization saved to: {output_path}")

def find_clade_in_forest(
	trees: list[Bio.Phylo.BaseTree.Tree],
	name: str | None = None,
	glottocode: str | None = None
) -> Bio.Phylo.BaseTree.Clade:
	"""
	Searches all trees for a specific node matching name or glottocode.
	"""
	for tree in trees:
		for clade in tree.find_clades():
			gl = GlottoLang.from_string(clade.name)
			
			if glottocode and gl.glottocode == glottocode:
				return clade
			if name and gl.name.casefold() == name.casefold():
				return clade
	raise ValueError(f'Could not find node matching "{glottocode or name}" in any tree.')

def _normalize_search_params(tree_or_clade: Bio.Phylo.BaseTree.Tree | Bio.Phylo.BaseTree.Clade) -> tuple[str, str]:
	gl = GlottoLang.from_string(tree_or_clade.name)
	name = gl.name
	glottocode = gl.glottocode
	return name, glottocode

def _clean_node_names(tree_mixin: Bio.Phylo.BaseTree.TreeMixin):
	"""
	Cleans labels for all nodes
	"""
	for clade in tree_mixin.find_clades():
		gl = GlottoLang.from_string(clade.name)
		clade.name = gl.name

if __name__ == "__main__":
	DEFAULT_TREE_FILE = "tree_glottolog_newick.txt"
	
	parser = argparse.ArgumentParser(
		description="Glottolog Tree Visualizer"
	)
	parser.add_argument("--name", help="Language name")
	parser.add_argument("--glottocode", help="Glottocode")
	parser.add_argument("--tree-file", default=DEFAULT_TREE_FILE, help="Path to Glottolog Newick tree file")
	
	parser.add_argument(
		"--tree-layout",
		choices=typing.get_args(dendroviz.models.TreeLayout),
		default="horizontal",
		help="Tree layout style"
	)
	parser.add_argument(
		"--line-style",
		choices=typing.get_args(dendroviz.models.LineStyle),
		default="straight",
		help="Line style"
	)
	
	args = parser.parse_args()
	
	# resolve tree file path: CLI > Default
	tree_file_path = args.tree_file
	tree_file_path = pathlib.Path(tree_file_path)
	
	if not tree_file_path.exists():
		print(f"Error: {tree_file_path} not found locally. Download {DEFAULT_TREE_FILE} from Glottolog website.")
		webbrowser.open("https://glottolog.org/download")
		exit(1)
	
	# resolve search params: CLI > Prompt
	name = args.name
	glottocode = args.glottocode
	if not (name or glottocode):
		name = input("Enter language name (or Enter to skip): ")
		if not name:
			glottocode = input("Enter Glottocode: ")
	
	# kwargs for visualization
	tree_gen_kwargs = {
		"tree_layout": args.tree_layout,
		"line_style":  args.line_style
	}
	
	draw_language_tree(tree_file_path, name=name, glottocode=glottocode, **tree_gen_kwargs)