formatitko/transform.py

from panflute import *
import re

# Import local files
from whitespace import *
from command import *
from util import *
from context import *
from group import Group

class FQuoted(Quoted):
	def __init__(self, *args, **kwargs):
		self.style = kwargs["style"]
		del kwargs["style"]
		super().__init__(*args, **kwargs)

def transform(e: Element, c: Context) -> Element: # Returns next sibling element to transform
	"""Transform the AST, making format-agnostic changes."""
	if isinstance(e, Whitespace) and bavlna(e, c):
		e = NBSP()

	if hasattr(e, "attributes"):
		# `if` attribute. Only show this element if flag is set.
		if "if" in e.attributes:
			if not c.is_flag_set(e.attributes["if"]):
				return nullify(e)
		# `ifn` attribute. Only show this element if flag is NOT set
		if "ifn" in e.attributes:
			if c.is_flag_set(e.attributes["ifn"]):
				return nullify(e)

		# `c` attribute. Execute a command with the name saved in this attribute.
		if (isinstance(e, Div) or isinstance(e, Span)) and "c" in e.attributes:
			if isinstance(e, Div):
				e = BlockCommand(*e.content, identifier=e.identifier, classes=e.classes, attributes=e.attributes)
			else:
				e = InlineCommand(*e.content, identifier=e.identifier, classes=e.classes, attributes=e.attributes)

		# `partial` attribute.
		# This is for including content from files with their own flags and
		# commands without affecting the state of the current document.
		if (isinstance(e, Div)) and "partial" in e.attributes:
			includedDoc = import_md(open(c.dir + "/" + e.attributes["partial"], "r").read())
			nContext = Context(includedDoc, e.attributes["partial"], c)
			language = includedDoc.get_metadata("language")
			includedDoc = includedDoc.walk(transform, nContext)
			e = Group(*includedDoc.content, metadata={"language": language})


	if isinstance(e, Quoted):
		quote_styles = {
			"cs": "cs",
			"en": "en",
			"sk": "cs",
			None: None
		}
		e = FQuoted(*e.content, quote_type=e.quote_type, style=quote_styles[c.get_metadata("language")])

	if isinstance(e, Image):
		e.attributes["source_dir"] = c.dir
		if not "no-srcset" in e.attributes:
			e.attributes["no-srcset"] = c.get_metadata("no-srcset") if c.get_metadata("no-srcset") is not None else False

	# Execute python code inside source code block
	if isinstance(e, CodeBlock) and hasattr(e, "classes") and "python" in e.classes and "run" in e.classes:
		e = Div(*executeCommand(e.text, None, c))
		e = e.walk(transform, c)

	## Command defines
	# possible TODO: def/longdef?
	if isinstance(e, CodeBlock) and hasattr(e, "classes") and "python" in e.classes and hasattr(e, "attributes")\
		and ("define" in e.attributes or "redefine" in e.attributes):
		e = handle_command_define(e, c)

	# Pass down metadata 'highlight' and 'highlight_style' as attribute to CodeBlocks
	if isinstance(e, CodeBlock):
		if not "highlight" in e.attributes:
			e.attributes["highlight"] = c.get_metadata("highlight") if c.get_metadata("highlight") is not None else True
		if not "style" in e.attributes:
			e.attributes["style"] = c.get_metadata("highlight_style") if c.get_metadata("highlight_style") is not None else "default"

	## Shorthands
	if isinstance(e, Span) and 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:]})

		## Handle import [#path/file.md]{}
		# This is the exact opposite of include. We take the commands
		# and flags but drop the content.
		elif re.match(r"^#.+$", e.content[0].text):
			importedDoc = import_md(open(c.dir + "/" + e.content[0].text[1:], "r").read())
			importedDoc.walk(transform, c)
			return nullify(e)

		## Handle metadata print [$something.something]{}
		elif re.match(r"^\$[\w.]+$", e.content[0].text):
			val = c.get_metadata(e.content[0].text[1:], False)
			if isinstance(val, MetaInlines):
				e = Span(*val.content)
				e = e.walk(transform, c)
			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)}'")

	## Execute commands
	# panflute's walk transforms the children first, then the root element, so
	# the content of the element the command receives is already transformed.
	if isinstance(e, Command):
		if not c.get_command(e.attributes["c"]):
			raise NameError(f"Command not defined '{e.attributes['c']}'.")
		e = e.replaceSelf(executeCommand(c.get_command(e.attributes["c"]), e, c))
		e = e.walk(transform, c)

	return e