formatitko/src/formatitko/transform_processor.py

from panflute import Element, ListContainer, Inline, Block
from panflute import Cite, Code, Emph, Image, LineBreak, Link, Math, Note, Quoted, RawInline, SmallCaps, SoftBreak, Space, Span, Str, Strikeout, Strong, Subscript, Superscript, Underline
from panflute import BlockQuote, BulletList, Citation, CodeBlock, Definition, DefinitionItem, DefinitionList, Div, Figure, Header, HorizontalRule, LineBlock, LineItem, ListItem, MetaBlocks, MetaBool, MetaInlines, MetaList, MetaMap, MetaString, Null, OrderedList, Para, Plain, RawBlock, Table, TableBody, TableFoot, TableHead
from panflute import TableRow, TableCell, Caption, Doc
from panflute import MetaValue
from typing import Union, Callable
from types import ModuleType

import os
import re
import warnings
import importlib
import json

from .whitespace import NBSP
from .elements import FQuoted
from .context import Group, InlineGroup, BlockGroup
from .util import nullify, import_md
from .context import Context, CommandCallable
from .whitespace import Whitespace, bavlna
from .command import BlockCommand, InlineCommand, CodeCommand, Command
from .command_util import handle_command_define, parse_command

ELCl = Union[Element, ListContainer, list[Union[Element, ListContainer]]]

class UnknownElementError(Exception):
	"An unknown Element has been passed to the TransformProcessor, probably because panflute introduced a new one."
	pass

class DoubleDocError(Exception):
	"TransformProcessor should only ever see a single Doc."
	pass

class TransformProcessor:
	
	context: Union[Context, None] = None
	root_file_path: str
	root_highlight_style: str = "default"
	_command_modules: list[tuple[Union[dict[str, CommandCallable], ModuleType], str]] = []
	TYPE_DICT: dict[type, Callable]

	def __init__(self, root_file_path: str):
		self.root_file_path = root_file_path

		self.TYPE_DICT = {
			TableRow: self.transform_TableRow,
			TableCell: self.transform_TableCell,
			Caption: self.transform_Caption,
			Doc: self.transform_Doc,
			LineItem: self.transform_LineItem,
			ListItem: self.transform_ListItem,

			BlockQuote: self.transform_BlockQuote,
			BulletList: self.transform_BulletList,
			Citation: self.transform_Citation,
			CodeBlock: self.transform_CodeBlock,
			Definition: self.transform_Definition,
			DefinitionItem: self.transform_DefinitionItem,
			DefinitionList: self.transform_DefinitionList,
			Div: self.transform_Div,
			Figure: self.transform_Figure,
			Header: self.transform_Header,
			HorizontalRule: self.transform_HorizontalRule,
			LineBlock: self.transform_LineBlock,
			MetaBlocks: self.transform_MetaBlocks,
			MetaBool: self.transform_MetaBool,
			MetaInlines: self.transform_MetaInlines,
			MetaList: self.transform_MetaList,
			MetaMap: self.transform_MetaMap,
			MetaString: self.transform_MetaString,
			Null: self.transform_Null,
			OrderedList: self.transform_OrderedList,
			Para: self.transform_Para,
			Plain: self.transform_Plain,
			RawBlock: self.transform_RawBlock,
			Table: self.transform_Table,
			TableBody: self.transform_TableBody,
			TableFoot: self.transform_TableFoot,
			TableHead: self.transform_TableHead,
			Group: self.transform_Group,
			InlineGroup: self.transform_InlineGroup,
			BlockGroup: self.transform_BlockGroup,

			Cite: self.transform_Cite,
			Code: self.transform_Code,
			Emph: self.transform_Emph,
			Image: self.transform_Image,
			LineBreak: self.transform_LineBreak,
			Link: self.transform_Link,
			Math: self.transform_Math,
			Note: self.transform_Note,
			Quoted: self.transform_Quoted,
			RawInline: self.transform_RawInline,
			SmallCaps: self.transform_SmallCaps,
			SoftBreak: self.transform_SoftBreak,
			Space: self.transform_Space,
			Span: self.transform_Span,
			Str: self.transform_Str,
			Strikeout: self.transform_Strikeout,
			Strong: self.transform_Strong,
			Subscript: self.transform_Subscript,
			Superscript: self.transform_Superscript,
			Underline: self.transform_Underline,
			NBSP: self.transform_NBSP,
			FQuoted: self.transform_FQuoted,

			InlineCommand: self.transform_InlineCommand,
			BlockCommand: self.transform_BlockCommand,
			CodeCommand: self.transform_CodeCommand
		}

	def add_command_module(self, module: Union[dict[str, CommandCallable], ModuleType], module_name: str=""):
		self._command_modules.append((module, module_name))

	def get_pretransformers(self) -> list[Callable[[ELCl],ELCl]]:
		return [self.handle_if_attribute, self.handle_ifnot_attribute]

	def get_posttransformers(self) -> list[Callable[[ELCl],ELCl]]:
		return []

	def transform(self, e: ELCl) -> ELCl:
		if isinstance(e, list):
			return self.transform_list(e)
		elif isinstance(e, ListContainer):
			return self.transform_ListContainer(e)

		for transformer in self.get_pretransformers():
			e = transformer(e)

		try:
			e = self.TYPE_DICT[type(e)](e)
		except KeyError:
			raise UnknownElementError(type(e))

		for transformer in self.get_posttransformers():
			e = transformer(e)

		return e


	def handle_if_attribute(self, e: ELCl) -> ELCl:
		# `if` attribute. Only show this element if flag is set.
		if hasattr(e, "attributes") and "if" in e.attributes:
			if not self.context.is_flag_set(e.attributes["if"]):
				return nullify(e)
		return e

	def handle_ifnot_attribute(self, e: ELCl) -> ELCl:
		# `ifnot` attribute. Only show this element if flag is NOT set
		if hasattr(e, "attributes") and "ifnot" in e.attributes:
			if self.context.is_flag_set(e.attributes["ifnot"]):
				return nullify(e)
		return e


	def transform_list(self, e: list[Union[Element, ListContainer]]) -> list[Union[Element, ListContainer]]:
		for i in range(len(e)):
			e[i] = self.transform(e[i])
		return e

	def transform_ListContainer(self, e: ListContainer) -> ListContainer:
		for i in range(len(e)):
			e[i] = self.transform(e[i])
		return e

	
	def transform_TableRow(self, e: TableRow) -> TableRow:
		e.content = self.transform(e.content)
		return e

	def transform_TableCell(self, e: TableCell) -> TableCell:
		e.content = self.transform(e.content)
		return e

	def transform_Caption(self, e: Caption) -> Caption:
		e.content = self.transform(e.content)
		return e

	def transform_LineItem(self, e: LineItem) -> LineItem:
		e.content = self.transform(e.content)
		return e

	def transform_ListItem(self, e: ListItem) -> ListItem:
		e.content = self.transform(e.content)
		return e

	def transform_BlockQuote(self, e: BlockQuote) -> BlockQuote:
		e.content = self.transform(e.content)
		return e

	def transform_BulletList(self, e: BulletList) -> BulletList:
		e.content = self.transform(e.content)
		return e

	def transform_Citation(self, e: Citation) -> Citation:
		e.content = self.transform(e.content)
		return e

	def transform_Definition(self, e: Definition) -> Definition:
		e.content = self.transform(e.content)
		return e

	def transform_DefinitionItem(self, e: DefinitionItem) -> DefinitionItem:
		e.content = self.transform(e.content)
		return e

	def transform_DefinitionList(self, e: DefinitionList) -> DefinitionList:
		e.content = self.transform(e.content)
		return e

	def transform_Header(self, e: Header) -> Header:
		e.content = self.transform(e.content)
		return e

	def transform_LineBlock(self, e: LineBlock) -> LineBlock:
		e.content = self.transform(e.content)
		return e

	def transform_MetaBlocks(self, e: MetaBlocks) -> MetaBlocks:
		e.content = self.transform(e.content)
		return e

	def transform_MetaBool(self, e: MetaBool) -> MetaBool:
		e.content = self.transform(e.content)
		return e

	def transform_MetaInlines(self, e: MetaInlines) -> MetaInlines:
		e.content = self.transform(e.content)
		return e

	def transform_MetaList(self, e: MetaList) -> MetaList:
		e.content = self.transform(e.content)
		return e

	def transform_MetaMap(self, e: MetaMap) -> MetaMap:
		e.content = self.transform(e.content)
		return e

	def transform_MetaString(self, e: MetaString) -> MetaString:
		e.content = self.transform(e.content)
		return e

	def transform_OrderedList(self, e: OrderedList) -> OrderedList:
		e.content = self.transform(e.content)
		return e

	def transform_Para(self, e: Para) -> Para:
		e.content = self.transform(e.content)
		return e

	def transform_Plain(self, e: Plain) -> Plain:
		e.content = self.transform(e.content)
		return e

	def transform_TableBody(self, e: TableBody) -> TableBody:
		e.content = self.transform(e.content)
		return e

	def transform_TableFoot(self, e: TableFoot) -> TableFoot:
		e.content = self.transform(e.content)
		return e

	def transform_TableHead(self, e: TableHead) -> TableHead:
		e.content = self.transform(e.content)
		return e

	def transform_Group(self, e: Group) -> Group:
		e.content = self.transform(e.content)
		return e

	def transform_InlineGroup(self, e: InlineGroup) -> InlineGroup:
		e.content = self.transform(e.content)
		return e

	def transform_BlockGroup(self, e: BlockGroup) -> BlockGroup:
		e.content = self.transform(e.content)
		return e

	def transform_Cite(self, e: Cite) -> Cite:
		e.content = self.transform(e.content)
		return e

	def transform_Emph(self, e: Emph) -> Emph:
		e.content = self.transform(e.content)
		return e

	def transform_Link(self, e: Link) -> Link:
		e.content = self.transform(e.content)
		return e

	def transform_Note(self, e: Note) -> Note:
		e.content = self.transform(e.content)
		return e

	def transform_SmallCaps(self, e: SmallCaps) -> SmallCaps:
		e.content = self.transform(e.content)
		return e

	def transform_Strikeout(self, e: Strikeout) -> Strikeout:
		e.content = self.transform(e.content)
		return e

	def transform_Strong(self, e: Strong) -> Strong:
		e.content = self.transform(e.content)
		return e

	def transform_Subscript(self, e: Subscript) -> Subscript:
		e.content = self.transform(e.content)
		return e

	def transform_Superscript(self, e: Superscript) -> Superscript:
		e.content = self.transform(e.content)
		return e

	def transform_Underline(self, e: Underline) -> Underline:
		e.content = self.transform(e.content)
		return e

	def transform_FQuoted(self, e: FQuoted) -> FQuoted:
		e.content = self.transform(e.content)
		return e

	def transform_Figure(self, e: Figure) -> Figure:
		e.content = self.transform(e.content)
		e.caption = self.transform(e.caption)
		return e

	def transform_Table(self, e: Table) -> Table:
		e.head = self.transform(e.head)
		e.content = self.transform(e.content)
		e.foot = self.transform(e.foot)
		return e


	def transform_Doc(self, e: Doc) -> Doc:
		if self.context is not None:
			raise DoubleDocError()
		self.context = Context(e, self.root_file_path)
		for module, module_name in self._command_modules:
			self.context.add_commands_from_module(module, module_name)
		e.content = self.transform(e.content)
		e.content = [BlockGroup(*e.content, context=self.context)]
		return e


	def transform_Quoted(self, e: Quoted) -> FQuoted:
		e.content = self.transform(e.content)
		quote_styles = {
			"cs": "cs",
			"en": "en",
			"sk": "cs",
			None: None
		}
		return FQuoted(*e.content, quote_type=e.quote_type, style=quote_styles[self.context.get_metadata("lang")])

	def transform_Image(self, e: Image) -> Image:
		e.content = self.transform(e.content)
		# OG now has Context so this is not needed per se, but I'm keeping this here for the handling of attribute > context > default value
		# Pass down "no-srcset" metadatum as attribute down to images.
		if not "no-srcset" in e.attributes:
			e.attributes["no-srcset"] = self.context.get_metadata("no-srcset") if self.context.get_metadata("no-srcset") is not None else False
		return e

	def create_Group(self, *content, new_context: Context, inline: bool=False) -> Group:
		old_context = self.context
		self.context = new_context
		content = self.transform([*content])
		self.context = old_context
		if inline:
			return InlineGroup(*content, context=new_context)
		else:
			return BlockGroup(*content, context=new_context)

	def transform_Div(self, e: Div) -> Union[Div, Group, Null, RawBlock]:
		e.content = self.transform(e.content)

		if "group" in e.classes:
			# `.group` class for Divs
			# Content of Div is enclosed in a separate context, all attributes are passed as metadata
			new_context = Context(Doc(), self.context.path, self.context, trusted=self.context.trusted)
			for attribute, value in e.attributes.items():
				new_context.set_metadata(attribute, value)
			return self.create_Group(*e.content, new_context=new_context)

		if "c" in e.attributes:
			# Commands can be called multiple ways, this handles the following syntax:
			# :::{c=commandname}
			# :::
			e = BlockCommand(*e.content, identifier=e.identifier, classes=e.classes, attributes=e.attributes)
			return self.transform(e)

		if "partial" in e.attributes:
			# `partial` attribute
			# Used to include a file which is executed in a different (child) Context.
			if not "type" in e.attributes:
				e.attributes["type"] = "md"
			if not self.context.trusted: # If we're in an untrusted context, we shouldn't allow inclusion of files outside the PWD.
				full_path = os.path.abspath(self.context.dir + "/" + e.attributes["partial"])
				pwd = os.path.abspath(".")
				if os.path.commonpath([full_path, pwd]) != os.path.commonpath([pwd]):
					return nullify(e)
			text = open(self.context.dir + "/" + e.attributes["partial"], "r").read()
			path = self.context.dir + "/" + e.attributes["partial"]
			if e.attributes["type"] == "md":
				includedDoc = import_md(text)
				trusted = True
				if "untrusted" in e.attributes and (e.attributes["untrusted"] == True or e.attributes["untrusted"] == 'True'):
					trusted = False
				if not self.context.trusted:
					trusted = False
				return self.create_Group(*includedDoc.content, new_context=Context(includedDoc, path, self.context, trusted=trusted))
			elif e.attributes["type"] in ["tex", "html"]:
				return RawBlock(text, e.attributes["type"])

		if "header_content" in e.classes:
			header_content = self.context.get_metadata("header_content")
			header_content = [] if header_content is None else header_content
			header_content.append(MetaBlocks(*self.transform(e.content)))
			self.context.set_metadata("header_content", header_content)
			return Null()

		if "footer_content" in e.classes:
			footer_content = self.context.get_metadata("footer_content")
			footer_content = [] if footer_content is None else footer_content
			footer_content.append(MetaBlocks(*self.transform(e.content)))
			self.context.set_metadata("footer_content", footer_content)
			return Null()

		if "lang" in e.attributes:
			warnings.warn("To set language in a way formátítko will understand, this Div has to have the `.group` class and be a Group.", UserWarning)

		return e

	def transform_Span(self, e: Span) -> Span:
		e.content = self.transform(e.content)

		if "group" in e.classes:
			# `.group` class for Spans
			# Content of Span is enclosed in a separate context, all attributes are passed as metadata
			new_context = Context(Doc(), self.context.path, self.context, trusted=self.context.trusted)
			for attribute, value in e.attributes.items():
				new_context.set_metadata(attribute, value)
			return self.create_Group(*e.content, new_context=new_context, inline=True)

		if "c" in e.attributes:
			# Commands can be called multiple ways, this handles the following syntax:
			# []{c=commandname} and
			e = InlineCommand(*e.content, identifier=e.identifier, classes=e.classes, attributes=e.attributes)
			return self.transform(e)

		if len(e.content) == 1 and isinstance(e.content[0], Str):
			## Handle special command shorthand [!commandname]{}
			if re.match(r"^![\w.]+$", e.content[0].text):
				e = InlineCommand(identifier=e.identifier, classes=e.classes, attributes={**e.attributes, "c": e.content[0].text[1:]})
				return self.transform(e)

			## Handle import [#ksp_formatitko as ksp]{}, [#ksp_formatitko]{type=module} or [#path/file.md]{type=md}
			# Import a python module as commands (type=module, the default) or
			# import all metadata from a md file, dropping its contents.
			elif re.match(r"^#.+$", e.content[0].text):
				if not "type" in e.attributes:
					e.attributes["type"] = "module"
				if e.attributes["type"] == "md":
					importedDoc = import_md(open(self.context.dir + "/" + e.content[0].text[1:], "r").read())
					self.transform(importedDoc.content)
				elif e.attributes["type"] == "module":
					matches = re.match(r"^(\w+)(?: as (\w+))?$", e.content[0].text[1:])
					if not matches:
						raise SyntaxError(f"`{e.content[0].text[1:]}`: invalid syntax")
					module = importlib.import_module(matches.group(1))
					module_name = matches.group(1) if matches.group(2) is None else matches.group(2)
					self.context.add_commands_from_module(module, module_name)
				elif e.attributes["type"] == "metadata":
					data = json.load(open(self.context.dir + "/" + e.content[0].text[1:], "r"))
					key = "" if not "key" in e.attributes else e.attributes["key"]
					self.context.import_metadata(data, key)
				else:
					raise SyntaxError(f"`{e.attributes['type']}`: invalid import type")

				return nullify(e)

			## Handle metadata print [$key1.key2]{}
			# This is a shorthand for just printing the content of some metadata.
			elif re.match(r"^\$[\w.]+$", e.content[0].text):
				val = self.context.get_metadata(e.content[0].text[1:], False)
				if isinstance(val, MetaInlines):
					e = Span(*val.content)
					e = self.transform(e)
				elif isinstance(val, MetaString):
					e = Span(Str(val.string))
				elif isinstance(val, MetaBool):
					e = Span(Str(str(val.boolean)))
				else:
					raise TypeError(f"Cannot print value of metadatum '{e.content[0].text[1:]}' of type '{type(val)}'")
				return e

		return e

	def transform_CodeBlock(self, e: CodeBlock) -> Union[CodeBlock, Div, Null]:
		if "markdown" in e.classes and "group" in e.classes:
			includedDoc = import_md(e.text)
			return self.create_Group(*includedDoc.content, new_context=Context(includedDoc, self.context.path, self.context, self.context.trusted))

		if "python" in e.classes and "run" in e.classes:
			if not self.context.trusted:
				return nullify(e)
			command_output = parse_command(e.text)(BlockCommand(), self.context)
			e = BlockCommand().replaceSelf(*([] if command_output is None else command_output))
			return self.transform(e)

		if "python" in e.classes and ("define" in e.attributes or "redefine" in e.attributes):
			if not self.context.trusted:
				return nullify(e)
			return handle_command_define(e, self.context)

		if "c" in e.attributes:
			return self.transform(CodeCommand(e.text, identifier=e.identifier, classes=e.classes, attributes=e.attributes))

		# Pass down metadata 'highlight' and 'highlight_style' as attribute to CodeBlocks
		# OG now has Context so this is not needed per se, but I'm keeping this here for the handling of attribute > context > default value
		if not "highlight" in e.attributes:
			e.attributes["highlight"] = self.context.get_metadata("highlight") if self.context.get_metadata("highlight") is not None else True
		if not "style" in e.attributes:
			e.attributes["style"] = self.context.get_metadata("highlight-style") if self.context.get_metadata("highlight-style") is not None else "default"
		return e

	def transform_Command(self, e: Command) -> Union[Div, Span]:
		if not self.context.get_command(e.attributes["c"]):
			raise NameError(f"Command not defined '{e.attributes['c']}'.")
		command_output = self.context.get_command(e.attributes["c"])(e, self.context)
		e = e.replaceSelf(*([] if command_output is None else command_output))
		return self.transform(e)

	def transform_InlineCommand(self, e: InlineCommand) -> Span:
		return self.transform_Command(e)

	def transform_BlockCommand(self, e: BlockCommand) -> Div:
		return self.transform_Command(e)

	def transform_CodeCommand(self, e: CodeCommand) -> Div:
		return self.transform_Command(e)

	def transform_Whitespace(self, e: Whitespace) -> Whitespace:
		if bavlna(e, self.context):
			return NBSP()
		else:
			return e

	def transform_SoftBreak(self, e: SoftBreak) -> Whitespace:
		return self.transform_Whitespace(e)

	def transform_Space(self, e: Space) -> Whitespace:
		return self.transform_Whitespace(e)


	def transform_NBSP(self, e: NBSP) -> NBSP:
		return e

	def transform_Str(self, e: Str) -> Str:
		return e

	def transform_RawInline(self, e: RawInline) -> RawInline:
		return e

	def transform_Math(self, e: Math) -> Math:
		return e

	def transform_LineBreak(self, e: LineBreak) -> LineBreak:
		return e

	def transform_Code(self, e: Code) -> Code:
		return e

	def transform_RawBlock(self, e: RawBlock) -> RawBlock:
		return e

	def transform_Null(self, e: Null) -> Null:
		return e

	def transform_HorizontalRule(self, e: HorizontalRule) -> HorizontalRule:
		return e