git-wtpp

#!/usr/bin/env uv --quiet run --no-project --script --isolated --refresh --
# /// script
# # Docopt issues SyntaxWarning in Python 3.12
# requires-python = ">=3.11,<3.12"
# dependencies = [
#   "docopt >=0.6.2",
# ]
# ///

"""
Usage:
	{prog} (push|pull|compare) [options] <remote>

Options:
	--force  Override worktree modifications
"""

import sys, locale, asyncio, subprocess, os, re, shlex, urllib.parse
import docopt

#TODO Decompose remote.<name>.wtpp-ssh-command to a reusable part and repo-specific part.
#TODO Deprecate remote.<name>.wtpp-ssh-host in favor of parsing the hostname from url.

async def main(*, args, prog):
	locale.setlocale(locale.LC_ALL, "")

	params = docopt.docopt(
		__doc__.replace("\t", " " * 4).format(prog=os.path.basename(prog)),
		argv=args,
		help=True,
		version=True,
		options_first=False
	)
	verbs = ["push", "pull", "compare"]
	verb = [i for i, v in enumerate(verbs) if params.pop(v)]
	assert len(verb) == 1, verb
	verb = verbs[verb[0]]
	del verbs
	remote_name = params.pop("<remote>")
	shall_force = params.pop("--force")
	assert not params, params

	log = asyncio.Queue()
	pending_tasks = []
	printer_task = asyncio.create_task(
		printer(queue=log, fo=sys.stdout)
	)

	#
	# ---
	#
	git_local = Git(cmd=("git",))
	await git_local.verify_version()
	git_remote = await git_local.create_remote_git(remote=remote_name)

	head_local = await git_local.get_rev_sha("HEAD")
	head_remote = await git_remote.get_rev_sha("HEAD")
	assert head_local == head_remote, (head_local, head_remote)
	await log.put(["Both local and remote are at commit ", head_local])

	if verb in ("push", "pull", "compare"):
		if local_stash := await git_local.stash_create():
			await log.put(("Created a local stash  ", local_stash))
		else:
			await log.put("Nothing to stash.")

		if remote_stash := await git_remote.stash_create():
			await log.put(["Created a remote stash ", remote_stash])
		else:
			await log.put("Nothing to stash.")

		if remote_stash != local_stash:
			remote_stash_ref = None

			if verb == "push":
				if shall_force:
					await log.put("Clearing remote worktree.")
					await git_remote.command("restore", "-s", "HEAD", "-W", "-S", ".", capture_output=False)
				if local_stash:
					remote_stash_ref = f"refs/wtpp/stashes/{local_stash}"
					await log.put(["Pushing the newly created stash to ", remote_name, " as ", remote_stash_ref])
					await git_local.push("--no-verify", "--force", remote_name, f"{local_stash}:{remote_stash_ref}")
					await log.put("Applying the stash on remote side.")
					await git_remote.stash_apply(local_stash)
			elif verb == "pull":
				if shall_force:
					await log.put("Clearing local worktree.")
					await git_local.command("restore", "-s", "HEAD", "-W", "-S", ".", capture_output=False)
				if remote_stash:
					#TODO Instead of creating a named reference, use something like: git fetch --upload-pack 'git -c uploadpack.allowAnySHA1InWant=true upload-pack' ssh://remote/~/repo af386f59e39f3f4d5606a25a341d63a09de55b73
					remote_stash_ref = f"refs/wtpp/stashes/{remote_stash}"
					await git_remote.command("update-ref", remote_stash_ref, remote_stash)
					await log.put("Fetching the newly created stash.")
					await git_local.fetch(remote_name, f"{remote_stash}")
					await log.put("Applying the stash on local side.")
					await git_local.stash_apply(remote_stash)
			elif verb == "compare":
				pass
			else:
				assert False, verb

			if remote_stash_ref:
				await log.put("Cleaning up references.")
				await git_remote.command("update-ref", "-d", remote_stash_ref)

		else:
			await log.put("Local and remote worktrees are identical.")
	else:
		assert False, verb
	#
	# ---
	#

	await asyncio.gather(*pending_tasks)
	current_task = asyncio.current_task()
	all_tasks = asyncio.all_tasks()
	assert {current_task, printer_task} == all_tasks, (current_task, all_tasks)

	await log.put(None)
	await printer_task


class Git(object):
	def __init__(self, *, cmd=None, host=None, user=None, cwd=None, stdout_encoding="UTF-8"):
		if cmd is None:
			cmd = ("git",)
		self._host = host
		self._user = user
		self._cmd = cmd
		self._cwd = cwd
		self._stdout_encoding = stdout_encoding
		self._is_windows = False

	async def create_remote_git(self, *, remote):
		remote_url = await self.config_get(f"remote.{remote}.url")

		host = None
		user = None
		repo_path = None

		if m := re.match(r"^(?P<username>[^@]+)@(?P<hostname>[^:]+):(?P<path>.*)$", remote_url):
			host = m["hostname"]
			user = m["username"]
			repo_path = m["path"]
		else:
			remote_url = urllib.parse.urlsplit(remote_url)
			assert remote_url.scheme == "ssh", remote_url
			assert remote_url.netloc == remote_url.hostname, remote_url
			host = remote_url.hostname
			repo_path = remote_url.path
			if repo_path.startswith("/~/"):
				repo_path = repo_path[1:]
			elif m := re.match(r"^/([A-Za-z]):/(.*)$", repo_path):
				repo_path = m[1] + ":/" + m[2]

		remote_git_cmd = shlex.split(await self.config_get(f"wtpp.{host}.gitCmd") or "git")

		remote_git = self.__class__(
			host=host,
			user=user,
			cmd=list(remote_git_cmd),
			stdout_encoding=self._stdout_encoding
		)
		await remote_git.verify_version()
		remote_git._cwd = repo_path
		return remote_git

	async def verify_version(self):
		version = (await self.command("version", text_output=True)).rstrip("\n")
		version = version.removeprefix("git version ").split(" ")[0]
		version = version.split(".")
		if version[-2] == "windows":
			del version[-2:]
			self._is_windows = True
		version = list(map(int, version))
		assert version[0] == 2 and version[1] >= 30 or version[0] > 2, list(version)

	async def command(self, *args, capture_output=True, returncode_ok=None, text_output=False, env=None):
		return_returncode = True
		if returncode_ok is None:
			return_returncode = False
			returncode_ok = lambda x: x == 0
		if returncode_ok is True:
			returncode_ok = lambda x: True

		if self._host is None:
			cmd = self._cmd[0]
			cmd_args = list(self._cmd[1:])
			if env is not None:
				new_env = dict(os.environ)
				new_env.update(env)
				env = new_env
			if self._cwd:
				cmd_args.append("-C")
				cmd_args.append(self._cwd)
		else:
			cmd = "ssh"
			cmd_args = [self._host]
			if self._user:
				cmd_args += ["-l", self._user]
			for k, v in (env or {}).items():
				assert_escaped_for_shell(k)
				assert_escaped_for_shell(v)
				if self._is_windows:
					assert "\"" not in k and "\"" not in v, (k, v)
					cmd_args.append(f"set \"{k}={v}\" && ")
				else:
					cmd_args.append(f"{k}={v}")
			cmd_args.extend(self._cmd)
			if self._cwd:
				remote_cwd = self._cwd.replace(" ", "\\ ")
				assert_escaped_for_shell(remote_cwd)
				cmd_args.append("-C")
				cmd_args.append(remote_cwd)

		cmd_args.extend(args)

		p = await asyncio.create_subprocess_exec(
			cmd, *cmd_args,
			stdout=subprocess.PIPE if capture_output else None,
			stderr=subprocess.PIPE if capture_output else None,
			env=env
		)
		stdout, stderr = await p.communicate()
		assert returncode_ok(p.returncode) and not stderr, (p.returncode, stderr, (cmd, cmd_args))
		result = stdout.decode(self._stdout_encoding) if text_output else stdout
		return (p.returncode, result) if return_returncode else result

	async def rev_parse(self, *args, verify=False):
		if verify:
			args = ["--verify", *args]
		result = await self.command("rev-parse", *args, text_output=True)
		if verify:
			result = result.strip()
		return result

	async def rev_list(self, *args):
		return (await self.command("rev-list", *args, text_output=True))

	async def get_rev_sha(self, rev):
		o = (await self.rev_parse(rev)).splitlines()
		assert len(o) == 1, o
		return o[0]

	async def get_rev_fullname(self, rev):
		o = (await self.rev_parse("--symbolic-full-name", rev)).splitlines()
		assert len(o) == 1, o
		return o[0]

	async def symbolic_ref(self, name, *, set_to=None):
		args = []
		if set_to is not None:
			args = ["-m", "git-wtpp", name, set_to]
		else:
			args = [name]
		returncode, output = await self.command("symbolic-ref", "--quiet", *args, text_output=True, returncode_ok=True)
		if returncode == 0:
			return output.strip()
		elif returncode == 1:
			return None
		else:
			assert False, returncode

	async def config_get(self, name):
		returncode, output = await self.command("config", "--get", name, text_output=True, returncode_ok=True)
		if returncode == 0:
			return output.rstrip("\n")
		elif returncode == 1:
			return None
		else:
			assert False, (returncode, name)

	async def stash_create(self):
		head_symbolic_ref = await self.symbolic_ref("HEAD")
		await self.command("update-ref", "--no-deref", "HEAD", "HEAD")
		returncode, o = await self.command("stash", "create", returncode_ok=lambda x: x in (0, 1), env={
			"GIT_AUTHOR_NAME": "None",
			"GIT_AUTHOR_EMAIL": "None",
			"GIT_AUTHOR_DATE": "1970-01-01T00:00:00Z",
			"GIT_COMMITTER_NAME": "None",
			"GIT_COMMITTER_EMAIL": "None",
			"GIT_COMMITTER_DATE": "1970-01-01T00:00:00Z",
		})
		result = self.extract_sha(o) if o != b"" else None
		if head_symbolic_ref is not None:
			await self.symbolic_ref("HEAD", set_to=head_symbolic_ref)
		return result

	async def stash_apply(self, sha):
		await self.command("stash", "apply", "--index", sha, capture_output=False)

	async def fetch(self, *args):
		await self.command("fetch", *args, capture_output=False)

	async def push(self, *args):
		await self.command("push", *args, capture_output=False)

	@staticmethod
	def extract_sha(b):
		m = re.match(rb"^([0-9a-fA-F]{40})\n$", b)
		assert m, b
		return m.groups()[0].decode("ascii")


def assert_escaped_for_shell(x: str) -> None:
	parts = shlex.split(x)
	assert len(parts) == 1, [x, parts]


async def printer(*, queue, fo):
	while True:
		item = await queue.get()
		if item is None:
			break
		if isinstance(item, (list, tuple)):
			print(*item, file=fo, sep="")
		else:
			print(item, file=fo)


def smain(argv=None):
	if argv is None:
		argv = sys.argv

	try:
		if sys.platform == "win32":
			asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

		return asyncio.run(
			main(
				args=argv[1:],
				prog=argv[0]
			),
			debug=False,
		)
	except KeyboardInterrupt:
		print(file=sys.stderr)


if __name__ == "__main__":
	sys.exit(smain())