pyesearch/pyesearch.py

#!/usr/bin/env python3

import sys
import os
import subprocess
import shutil
from glob import glob
import re


def first(iterable, default=None):
	it = iter(iterable)
	try:
		return next(it)
	except StopIteration:
		return default


# Config reading

def read_ini(contents, collector=None):
	section = None

	if collector is None:
		collector = {}

	for line in map(str.strip, contents):
		line = re.sub(r"\s*#.*$", "", line)

		m = re.fullmatch(r"\[(.+)\]", line)
		if m:
			section = m.group(1)
			if section not in collector:
				collector[section] = {}
			continue

		m = re.fullmatch(r"([^=]+?)\s*=\s*(.*?)", line)
		if m:
			collector[section][m.group(1)] = m.group(2)
			continue

	return collector


def read_config(path):
	collector = {}
	# In Gentoo config files can be directories instead of files. The files in the directory are then
	# in evaluated alphabetical order.
	for repo_conf in sorted(glob(os.path.join(path, "*"))):
		with open(repo_conf) as fh:
			collector = read_ini(fh, collector)

	return collector


def repo_locations():
	result = []
	for name, defs in read_config("/etc/portage/repos.conf").items():
		if "location" in defs:
			result.append((name, defs["location"]))
	return result


def default_repo():
	return read_config("/etc/portage/repos.conf").get("DEFAULT", {}).get("main-repo", None)


# Searching

def list_from_print0(data):
	data = data.rstrip(b"\x00")

	if not data:
		return []

	return [
		x.decode()
		for x in data.split(b"\x00")
	]


FIND_CMD = shutil.which("find")
def find(query, paths, depth=None):
	depth_args = ["-mindepth", f"{depth}", "-maxdepth", f"{depth}"] if depth is not None else []

	proc = subprocess.run(
		[FIND_CMD, *paths, *depth_args, "-iname", f"*{query}*.ebuild", "-print0"],
		check=True, stdout=subprocess.PIPE
	)
	return list_from_print0(proc.stdout)


FD_CMD = shutil.which("fd")
def fd(query, paths, depth):
	depth_args = [f"--exact-depth={depth}"] if depth is not None else []

	proc = subprocess.run(
		[FD_CMD, *depth_args, "--print0", "--extension=ebuild", "--", query, *paths],
		check=True, stdout=subprocess.PIPE
	)
	return list_from_print0(proc.stdout)


if FD_CMD is not None:
	search_backend = fd
elif FIND_CMD is not None:
	search_backend = find
else:
	raise Exception("Neither `fd` nor `find` found in PATH")


# Ebuild reading

def description_from_ebuild(path):
	with open(path) as fh:
		for line in map(str.strip, fh):
			m = re.fullmatch(r'DESCRIPTION=(?:"(.*)"|[^"](.*))', line)
			if m:
				return m.group(1).replace(r'\"', '"') if m.group(1) else m.group(2)
	return None


def ebuild_info(path, repos):
	m = re.fullmatch(r"(.+)/([^/]+)/([^/]+)/\3-([^/]+)\.ebuild", path)
	assert m, path
	PN = f"{m.group(2)}/{m.group(3)}"
	PV = f"{m.group(4)}"

	return {
		"P": f"{PN}-{PV}",
		"PN": PN,
		"PV": PV,
		"category": m.group(2),
		"name": m.group(3),
		"repo": first((name for name, path in repos if path == m.group(1))),
		"repo_path": m.group(1),
		"ebuild_path": path,
	}


def ebuild_extra_info(info):
	return {
		**info,
		"description": description_from_ebuild(info["ebuild_path"]),
		"installed": os.path.exists(f"/var/db/pkg/{info['P']}"),
	}


# Formatting

def longest_category_and_nameversion(infos, default_repo=None):
	return (
		max(len(x["category"]) for x in infos),
		max(
			len(f"{x['name']}-{x['PV']}") +
			len(f"::{x['repo']}" if x["repo"] != default_repo else "")
		for x in infos)
	)


def format_item(info, longest_cat=0, longest_namever=0, default_repo=None):
	colour = "\x1b[0m" if info["installed"] else "\x1b[0;90m"
	name = info['name']
	name_colour = "\x1b[92m" if info["installed"] else "\x1b[36m"
	version = info['PV']
	repo = f"::{info['repo']}" if info["repo"] != default_repo else ""

	description_colour = "\x1b[0m"
	description = info["description"] if info["description"] else ""

	filler = (
		" " +
		("." * (longest_namever - len(name) - 1 - len(version) - len(repo)))
	) if info["description"] else ""
	return f"{colour}{info['category']:>{longest_cat}}/" \
		f"{name_colour}{name}{colour}-{version}{colour}{repo}{filler} " \
		f"{description_colour}{description}"



def main(argv):
	if len(argv) != 2 or not argv[1]:
		print("Need exactly 1 argument: search query", file=sys.stderr)
		return 1

	repos = repo_locations()
	default = default_repo()

	if not repos:
		raise Exception("Could not find any repositories in /etc/portage/repos.conf")

	data = search_backend(argv[1], (r[1] for r in repos), depth=3)
	if not data:
		print("No results")
		return 0

	data = [ebuild_info(x, repos=repos) for x in data]
	data.sort(key=lambda x: x["P"])

	longest_cat, longest_namever = longest_category_and_nameversion(data, default)

	for x in data:
		print(format_item(
			ebuild_extra_info(x),
			longest_cat=longest_cat, longest_namever=longest_namever, default_repo=default
		))

	return 0

if __name__ == "__main__":
	sys.exit(main(sys.argv))