formatitko/src/formatitko/context.py

from panflute import Doc, Element, Div, Span, Header, Figure

from typing import Union, Callable
from types import ModuleType
import os

from .command import Command

CommandCallable = Callable[[Command, 'Context', 'NOPProcessor'], list[Element]] # This is here because of a wild circular import dependency between many functions and classes

# This class is used to keep state while transforming the document using
# transform.py. For the context to be available to the html and TeX generators,
# individual keys must be manually assigned to the individual elements. This is
# done in transform.py.
#
# The context is also aware of its parent contexts and relevant data (such as
# metadata and commands) can be read from the closest parent context. Writing
# only happens to the current one.
#
# This class is basically an extension to panflute's doc, this is why metadata
# is read directly from it.
def default_section_number_generator(e: Header, context: 'Context') -> list[Union[str, int]]:
	l = e.level
	section_counters = context.get_data("section_counters")
	section_counters[l-1] += 1
	for i in range(l, len(section_counters)):
		section_counters[i] = 0
	return list(section_counters[:l])

def default_figure_number_generator(e: Figure, context: 'Context') -> Union[str, int]:
	figure_type = e.attributes.get("type", "img")
	figure_counters = context.get_data("figure_counters")
	figure_counters.setdefault(figure_type, 0)
	figure_counters[figure_type] += 1
	return figure_counters[figure_type]



class Context:
	parent: Union["Context", None]
	_commands: dict[str, Union[CommandCallable, None]]
	_data: dict[str, object]
	doc: Doc
	trusted: bool
	path: str
	dir: str
	filename: str
	root_dir: str  # Absolute path to the dir of the file formátítko was called on
	rel_dir: str   # Relative path to the current dir from the root dir
	deps: set[str]

	section_counters: list[int]
	number_generator: Callable[[Header, 'Context'], str]

	def __init__(self, doc: Doc, path: str, parent: Union['Context', None]=None, trusted: bool=True):
		self.parent = parent
		self._commands = {}
		self._data = {}
		self.doc = doc
		self.trusted = trusted
		self.path = path
		self.dir = os.path.dirname(path) if os.path.dirname(path) != "" else "."
		self.filename = os.path.basename(path)
		self.root_dir = parent.root_dir if parent else os.path.abspath(self.dir)
		self.rel_dir = os.path.relpath(self.dir, self.root_dir)
		self.deps = set()
		self.add_dep(path)
		if self.get_metadata("flags", immediate=True) is None:
			self.set_metadata("flags", {})
		if not parent:
			self.set_data('section_number_generator', default_section_number_generator)
			self.set_data('figure_number_generator', default_figure_number_generator)
			self.set_data('section_counters', [0 for i in range(6)])
			self.set_data('figure_counters', {})
			self.set_data('obj_map', {})

	def get_command(self, command: str) -> Union[CommandCallable, None]:
		if command in self._commands:
			return self._commands[command]
		elif self.parent:
			return self.parent.get_command(command)
		else:
			return None

	def set_command(self, command: str, val: CommandCallable):
		self._commands[command] = val

	def unset_command(self, command: str):
		del self._commands[command]

	def add_commands_from_module(self, module: Union[dict[str, CommandCallable], ModuleType], module_name: str=""):
		if isinstance(module, ModuleType):
			module = module.__dict__
		prefix = module_name+"." if module_name else ""
		for name, func in module.items():
			if isinstance(func, Callable):
				self.set_command(prefix+name, func)

	def is_flag_set(self, flag: str):
		if self.get_metadata("flags."+flag) is not None:
			if self.get_metadata("flags."+flag):
				return True
			else:
				return False
		elif self.parent:
			return self.parent.is_flag_set(flag)
		else:
			return False

	def set_flag(self, flag: str, val: bool):
		self.set_metadata("flags."+flag, val)

	def unset_flag(self, flag: str):
		self.unset_metadata("flags."+flag)

	def get_metadata(self, key: str, simple: bool=True, immediate: bool=False):
		value = self.doc.get_metadata(key, None, simple)
		if value is not None:
			return value
		elif self.parent and not immediate:
			return self.parent.get_metadata(key)
		else:
			return None

	def set_metadata(self, key: str, value):
		meta = self.doc.metadata
		keys = key.split(".")
		for k in keys[:-1]:
			meta = meta[k]
		meta[keys[-1]] = value

	def unset_metadata(self, key: str):
		meta = self.doc.metadata
		keys = key.split(".")
		for k in keys[:-1]:
			meta = meta[k]
		del meta.content[key[-1]] # A hack because MetaMap doesn't have a __delitem__
	
	def import_metadata(self, data, key: str=""):
		if isinstance(data, dict) and isinstance(self.get_metadata(key), dict):
			for subkey, value in enumerate(data):
				self.import_metadata(value, key+"."+subkey if key != "" else subkey)
		else:
			self.set_metadata(key, data)

	def get_data(self, key: str, immediate: bool=False):
		data = self._data
		keys = key.split(".")
		try:
			for k in keys:
				data = data[k]
			return data
		except KeyError:
			if self.parent and not immediate:
				return self.parent.get_data(key)
			else:
				return None

	def set_data(self, key: str, value: object):
		data = self._data
		keys = key.split(".")
		for k in keys[:-1]:
			try:
				data = data[k]
			except KeyError:
				data[k] = {}
				data = data[k]
		data[keys[-1]] = value

	def unset_data(self, key: str):
		if key == "":
			self._doc = {}
		data = self._data
		keys = key.split(".")
		for k in keys[:-1]:
			data = data[k]
		del data[keys[-1]]

	def get_deps(self) -> list[str]:
		if self.parent is not None:
			return self.parent.get_deps()
		else:
			return self.deps

	def add_dep(self, dep: str):
		self.get_deps().add(os.path.abspath(dep))

	def add_deps(self, deps: list[str]):
		self.get_deps().update([os.path.abspath(path) for path in deps])


def get_context_from_doc(doc: Doc) -> Context:
	if len(doc.content) == 1 and isinstance(doc.content[0], Group):
		return doc.content[0].context
	else:
		return None

# This is a custom element which creates \begingroup \endgroup groups in TeX
# and also causes KaTeX math blocks to be isolated in a similar way.
#
# Whenever a new context is created, its content should be eclosed in a group and vice-versa.
class Group(Element):
	metadata: dict
	context: Context
	def __init__(self, *args, context:Context, metadata={}, **kwargs):
		self.metadata = metadata # This is only here for backwards compatibility with old html.py, tex.py and transform.py. FIXME: Remove this when the time comes.
		self.context = context
		super().__init__(*args, **kwargs)

class BlockGroup(Group, Div):
	pass

class InlineGroup(Group, Span):
	pass