basgi/dev.py

#!/usr/bin/env python3
import asyncio
import os
import shlex
import subprocess
import sys
import time
import tomllib

from collections.abc import Callable
from datetime import datetime, timedelta
from pathlib import Path
from shutil import rmtree
from typing import TypedDict

try:
	import watchfiles

	from click import echo, group, option
	from tomlkit.items import Array, String, StringType, Trivia
	from tomlkit.toml_file import TOMLFile

except ImportError:
	print("Installing missing dependencies...")
	deps = " ".join(["build", "click", "tomlkit", "watchfiles"])
	subprocess.run(shlex.split(f"{sys.executable} -m pip install {deps}"))

	print("Restarting script...")
	subprocess.run([sys.executable, *sys.argv])
	sys.exit()


REPO = Path(__file__).resolve().parent
CLEAN_DIRS = ["build", "dist", "dist-pypi", "barkshark_asgi.egg-info", "docs/_build"]
IGNORE_DIRS = ["build", "dist", "dist-pypi", "docs", ".git", "barkshark_asgi.egg-info"]
IGNORE_PATHS = tuple(str(REPO.joinpath(path)) for path in IGNORE_DIRS)


class WatchfilesOptions(TypedDict):
	watch_filter: Callable[[watchfiles.Change, str], bool]
	recursive: bool
	ignore_permission_denied: bool
	rust_timeout: int


@group("cli")
def cli() -> None:
	...



@cli.command("clean")
def cli_clean() -> None:
	for directory in CLEAN_DIRS:
		try:
			rmtree(REPO.joinpath(directory))

		except FileNotFoundError:
			pass

	echo("Cleaned up build files")


@cli.command("install")
def cli_install_deps() -> None:
	with open("pyproject.toml", "rb") as fd:
		pyproject = tomllib.load(fd)

	dependencies = pyproject["project"]["dependencies"]
	dependencies.extend(pyproject["project"]["optional-dependencies"]["dev"])
	dependencies.extend(pyproject["project"]["optional-dependencies"]["docs"])
	dependencies = list(dep.replace(" ", "") for dep in dependencies)

	run_python("-m", "pip", "install", "-U", "pip", "setuptools", "wheel")
	run_python("-m", "pip", "install", *dependencies)
	echo("Installed dependencies :3")


@cli.command("lint")
@option("--path", "-p", type = Path, default = REPO.joinpath("basgi"))
@option("--watch", "-w", is_flag = True, help = "Watch for changes to the source")
def cli_lint(path: Path, watch: bool) -> None:
	path = path.expanduser().resolve()

	if watch:
		script = str(Path(__file__).resolve())
		handle_run_watcher(script, "lint", "--path", str(path))
		return

	echo("----- flake8 -----")
	run_python("-m", "flake8", str(path))

	echo("\n----- mypy -----")
	os.environ["MYPYPATH"] = str(REPO.joinpath("stubs"))
	run_python("-m", "mypy", str(path))


@cli.command("build")
def cli_build() -> None:
	cli_update_files.callback() # type: ignore
	run_python("-m", " build", "--outdir", "dist-pypi")


@cli.command("update-files")
def cli_update_files() -> None:
	project_file = TOMLFile(REPO.joinpath("pyproject.toml"))
	project = project_file.read()
	paths = []

	for path in REPO.joinpath("basgi").rglob("*"):
		if path.is_file() and not path.suffix == ".py":
			strpath = str(path.parent).split("basgi")[1][1:] + "/*"

			if strpath in paths:
				continue

			paths.append(strpath)

	parsed_paths = [String(StringType.SLB, path, path, Trivia()) for path in paths]
	files = Array(parsed_paths, multiline = True, trivia = Trivia(indent = "\t")) # type: ignore
	project["tool"]["setuptools"]["package-data"]["basgi"] = files # type: ignore
	project_file.write(project)


@cli.command("generate-stubs")
def cli_generate_stubs():
	subprocess.run(
		["-m", "mypy.stubgen", "-o", "stubs", "-p", "multipart", "--export-less"]
	)


def run_python(*arguments: str) -> subprocess.CompletedProcess[bytes]:
	return subprocess.run([sys.executable, *arguments])


def handle_run_watcher(*command: str) -> None:
	asyncio.run(_handle_run_watcher(*command))


async def _handle_run_watcher(*command: str) -> None:
	proc: subprocess.Popen[bytes] = subprocess.Popen([sys.executable, *command])
	last_restart: datetime = datetime.now()

	options: WatchfilesOptions = {
		"watch_filter": lambda _, path: path.endswith(".py"),
		"recursive": True,
		"ignore_permission_denied": True,
		"rust_timeout": 1000
	}

	async for changes in watchfiles.awatch(REPO.joinpath("basgi"), **options):
		skip = False

		for _, path in changes:
			if path.startswith(IGNORE_PATHS):
				skip = True

		if skip:
			continue


		if datetime.now() - timedelta(seconds = 3) < last_restart:
			continue

		if proc.poll() is None:
			echo(f"Terminating process {proc.pid}")
			proc.terminate()

			sec = 0.0

			while proc.poll() is None:
				time.sleep(0.1)
				sec += 0.1

				if sec < 5.0:
					continue

				echo("Failed to terminate. Killing process...")
				proc.kill()
				break

			echo("Process terminated")

		proc = subprocess.Popen([sys.executable, *command])
		last_restart = datetime.now()
		echo(f"Started processes with PID: {proc.pid}")


if __name__ == "__main__":
	cli()