formatitko/katex.py

import socket
import subprocess
import tempfile
import json
import os
from typing import Dict


class KatexError(Exception):
	pass

class KatexClient:
	def __init__(self):
		self._client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
		self._temp_dir = tempfile.TemporaryDirectory(prefix='formatitko')
		self._socket_file = self._temp_dir.name + "/katex-socket"
		self._server_process = subprocess.Popen(["node", "./katex-server/index.mjs", self._socket_file])
		while not os.path.exists(self._socket_file):
			pass
		while True:
			try:
				self._client.connect(self._socket_file)
			except ConnectionRefusedError:
				continue
			break

	def render(self, tex: str, options: Dict={}):
		options["globalGroup"] = True
		self._client.sendall((json.dumps({"formulas":[{"tex":tex}], "options":options})+"\n").encode("utf-8"))
		data = self._client.recv(1024)
		while data[-1] != 0x0a:
			data += self._client.recv(128)
		response = json.loads(data)
		if "error" in response:
			raise Exception(response["error"])
		if "error" in response["results"][0]:
			raise KatexError(response["results"][0]["error"])
		else:
			return response["results"][0]["html"]

	def begingroup(self):
		self._client.sendall("begingroup\n".encode("utf-8"))

	def endgroup(self):
		self._client.sendall("endgroup\n".encode("utf-8"))