basgi/basgi/template.py

from __future__ import annotations

import sass

from collections.abc import Callable
from hamlish_jinja import HamlishExtension, OutputMode
from jinja2 import Environment, FileSystemLoader
from jinja2.ext import Extension
from os.path import splitext
from pathlib import Path
from typing import Any
from xml.dom import minidom
from xml.etree import ElementTree

from .enums import SassOutputStyle
from .misc import Color


TemplateContextCallback = Callable[["Template", dict[str, Any]], dict[str, Any]]


class SassExtension(Extension):
	"An extension for Jinja2 that adds support for sass and scss compiling."


	def __init__(self, environment: Template):
		Extension.__init__(self, environment)

		self.output_style: SassOutputStyle = SassOutputStyle.NESTED
		"Output style to use when rendering"

		self.include_paths: list[str] = []
		"Paths to search when using the '@include' statement"

		self._exts: tuple[str, str] = (".sass", ".scss")


	def __repr__(self) -> str:
		return f"SassExtension(output_style='{self.output_style.value}')"


	def preprocess(self, source: str, name: str | None, filename: str | None = None) -> str:
		"""
			Transpile a sass or scss template into a jinja css template

			:param source: Full text source of the template
			:param name: Name of the template
			:param filename: Path to the template
			:raises CompileError: When the template cannot be parsed
		"""

		if (tpl_name := filename or name) is None:
			return source

		if (ext := splitext(tpl_name)[1]) not in self._exts:
			return source

		return sass.compile(
			string = source,
			output_style = self.output_style.value,
			indented = ext == ".sass"
		)


class Template(Environment):
	"Custom jinja environment with extensions to handle sass, scss, and haml"


	def __init__(self,
				*search: str | Path,
				context_function: TemplateContextCallback | None = None,
				**global_env: Any):
		"""
			Create a new template environment

			:param search: Paths for the environment to search for template files
			:param context_function: Function to call on the context to be passed to the template
				being rendered
			:param global_env: Variables to be included in the context on every render
		"""

		self.search: FileSystemLoader = FileSystemLoader([])

		super().__init__(
			loader = self.search,
			lstrip_blocks = True,
			trim_blocks = True,
			extensions = [
				HamlishExtension,
				SassExtension
			]
		)

		for path in search:
			self.add_search_path(path)

		self.autoescape: bool = True
		"Escape HTML output from variables and methods"

		self.context_function: TemplateContextCallback | None = None
		"Function to call on the context to be passed to the template being rendered"

		if context_function:
			self.set_context_function(context_function)

		self.global_env: dict[str, Any] = {
			"cleanhtml": lambda text: "".join(ElementTree.fromstring(text).itertext()),
			"color": Color,
			"lighten": lambda c, v: Color(c).lighten(v),
			"darken": lambda c, v: Color(c).darken(v),
			"saturate": lambda c, v: Color(c).saturate(v),
			"desaturate": lambda c, v: Color(c).desaturate(v),
			"rgba": lambda c, v: Color(c).rgba(v),
			**global_env
		}
		"Variables to be included in the context on every render"

		self.hamlish_file_extensions = (".haml", ".jhaml", ".jaml")
		self.hamlish_enable_div_shortcut = True
		self.hamlish_mode = OutputMode.INDENTED


	@property
	def sass(self) -> SassExtension:
		"Get the current ``SassExtension`` object"

		return self.extensions["basgi.template.SassExtension"] # type: ignore[return-value]


	def add_search_path(self, path: Path | str, index: int = 0) -> None:
		"""
			Add a path to be used when finding templates

			:param path: Path to be added
			:param index: Position to insert the path into the list. Gets inserted at the top by
				default
		"""

		if isinstance(path, str):
			path = Path(path).expanduser().resolve()

		if not path.exists():
			raise FileNotFoundError(f"Cannot find search path: {path}")

		if str(path) not in self.search.searchpath:
			self.search.searchpath.insert(index, str(path))
			self.sass.include_paths.append(str(path))


	def set_context_function(self, context: TemplateContextCallback) -> TemplateContextCallback:
		"""
			Set the function to be called on the context when rendering a template. Can be used as
			a decorator.

			:param context: Function to modify the context
		"""

		if not hasattr(context, "__call__"):
			raise TypeError("Context is not callable")

		if not isinstance(context(self, {}), dict):
			raise ValueError("Context does not return a dict object")

		self.context_function = context
		return context


	def add_filter(self, func: Callable[..., Any], name: str | None = None) -> None:
		"""
			Add a filter function. See
			`Jinja's docs <https://jinja.palletsprojects.com/en/3.1.x/templates/#filters>`_.

			:param func: Callable to be added as a filter
			:param name: Name to use when calling the filter. Defaults to the name of the function.
		"""

		if name is None:
			name = func.__name__

		self.filters[name] = func


	def render(self,
				template_name: str,
				context: dict[str, Any] | None = None,
				pprint: bool = False) -> str:
		"""
			Render a template file into a string

			:param template_name: Name of the template file to be rendered
			:param pprint: Format the resulting html/xml to be pretty
			:param context: Data to be passed to the template
		"""

		if context is None:
			context = {}

		context.update(self.global_env.copy())

		if self.context_function:
			context = self.context_function(self, context)

		result = self.get_template(template_name).render(context)

		if pprint and template_name.lower().endswith(("haml", "jhaml", "jaml", "html", "xml")):
			return minidom.parseString(result).toprettyxml(indent="  ")

		return result