find_and_open_file

#!/bin/python3

# complete rewrite of zzzfoo, to be simpler and work on my system

from recoll import recoll
import subprocess as sp
from subprocess import PIPE
import argparse
from dataclasses import dataclass
from tomllib import load as load_toml
import os
from sys import stderr
import typing

#
# ============================================================================================================
#

# maximum number of results to return
# making an assumption about how many results the user is willing to dig through
MAX_RESULTS: int = 20

FALLBACK_DEFAULT_COMMAND_FORMAT: list[str] = ['xdg-open', '{file}']

#
# config stuff ===============================================================================================
#

@dataclass(init=False)
class Config:
    pdf_command_format: list[str] | None
    default_command_format: list[str] | None

    def __init__(self):
        for member_name in type(self).__annotations__.keys():
            setattr(self, member_name, None)

    @staticmethod
    def get_member_type(member_name: str):
        member_type: type | None = Config.__annotations__.get(member_name)
        assert(member_type is not None)
        # assumes that all members in this class are written exactly as `some_type | None`, in that order.
        return typing.get_args(member_type)[0]


def get_config_file_location() -> str | None:
    # TODO check in directories https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html

    home_dir: str = os.path.expanduser('~')
    assert(home_dir != '~')

    config_path = home_dir + '/.config/fff.toml'
    if os.path.isfile(config_path): return config_path

    return None

def get_config() -> Config:

    validated_config = Config()

    config_file: str | None = get_config_file_location()
    if (config_file is None):
        print('No config file found')
        return validated_config

    print(f'Using config file `{config_file}`')

    with open(config_file, 'rb') as f:
        config: dict = load_toml(f)

    config_error: bool = False
    valid_config_keys = Config.__annotations__.keys()
    for key, val in config.items():

        if (key not in valid_config_keys):
            stderr.write(f'error: config file: invalid key `{key}`')
            config_error = True
            continue

        if (getattr(validated_config, key) is not None):
            stderr.write(f'error: config file: key `{key}` is set multiple times')
            config_error = True
            continue

        expected_type: type = Config.get_member_type(key)
        try: casted_val = expected_type(val)
        except(ValueError):
            stderr.write(f'error: config file: value of `{key}` must have type `{expected_type}`')
            config_error = True
            continue

        setattr(validated_config, key, casted_val)

    if (config_error):
        stderr.write('error: Aborting due to config errors.')
        exit(1)

    return validated_config

g_config: Config = get_config()

#
# ============================================================================================================
#

def get_pdf_command(fname: str, page: int | None) -> list[str] | None:

    if page is None: page = 0

    if (g_config.pdf_command_format is None): return None

    cmd: list[str] = g_config.pdf_command_format
    for i in range(len(cmd)):
        cmd[i] = cmd[i].replace('{file}', fname)
        cmd[i] = cmd[i].replace('{page}', str(page))
    return cmd

def get_generic_file_command(fname: str) -> list[str]:

    cmd: list[str]
    if (g_config.default_command_format is not None):
        cmd = g_config.default_command_format
    else:
        cmd = FALLBACK_DEFAULT_COMMAND_FORMAT

    for i in range(len(cmd)):
        cmd[i] = cmd[i].replace('{file}', fname)

    return cmd

def parse_args():
    parser = argparse.ArgumentParser(description='find files via a Rofi dialog')
    parser.add_argument('--fulltext',
        dest='full_text_search',
        action='store_const', const=True,
        default=False,
        help='also do a full text search within files'
    )
    parser.add_argument('--pdf',
        dest='pdf_only',
        action='store_const', const=True,
        default=False,
        help='only return `.pdf` files'
    )
    return parser.parse_args()

def prompt_search_query(fulltext: bool, pdf_only: bool):
    prompt = 'file search' if not fulltext else 'text search'
    if pdf_only: prompt = 'PDF ' + prompt

    ran = sp.run(['/bin/rofi', '-dmenu', '-p', prompt], stdout=PIPE, text=True)
    if ran.returncode != 0: raise RuntimeError('Rofi (search query dialog) returned non-zero')

    return ran.stdout

@dataclass
class SearchResult:
    file_url: str
    filename: str
    snippets: list[str]

def search(user_query: str, fulltext: bool, pdf_only: bool) -> list[SearchResult]:
    db = recoll.connect()
    dbquery = db.query()
    # if using sortby(), it must be before executing the search
    dbquery.sortby('relevancyrating', ascending=False)

    # build the query
    # `,` indicates AND
    query = []
    # fulltext or filename only?
    if fulltext: query.append(user_query)
    else: query.append('filename:' + user_query.replace(' ',','))
    #
    if pdf_only: query.append('ext:pdf')
    #
    query_onestring = ' '.join(query)

    # run the query
    _n_results = dbquery.execute(query_onestring)
    query_results = dbquery.fetchmany(MAX_RESULTS)

    results: list[SearchResult] = [
        SearchResult(
             file_url=result.url,
             filename=result.filename,
             snippets=dbquery.getsnippets(result, nohl=True) if fulltext else []
         )
        for result in query_results
    ]
    return results

@dataclass
class SelectableEntry:
    filename: str
    file_url: str
    page: int | None
    snippet: str | None

# returns the file url
def prompt_select_result(results: list[SearchResult], search_was_fulltext: bool) -> SelectableEntry:

    entries: list[SelectableEntry] = []

    idx = 0
    for result in results:
        if (search_was_fulltext):
            for snippet in result.snippets:
                entries.append(SelectableEntry(
                    filename=result.filename,
                    file_url=result.file_url,
                    # `snippet` is a tuple: (page number, matched word, actual snippet)
                    page=int(snippet[0]),
                    snippet=snippet[2]
                ))
        else:
            entries.append(SelectableEntry(
                filename=result.filename,
                file_url=result.file_url,
                page=None,
                snippet=None
            ))


    text_entries: list[str] = []
    idx = 0
    for entry in entries:
        text_entries.append(
            f'{idx} ' # prepend index to the filename
            + (f'<b>{entry.filename}</b>' if search_was_fulltext else f'{entry.filename}')
            + (f'\n<small>    p{entry.page}: {repr(entry.snippet)}</small>' if search_was_fulltext else '')
        )
        idx += 1

    entries_onestring = '\0'.join(text_entries)

    # make user select a file
    ran = sp.run(
        [
            '/bin/rofi',
            '-dmenu',
            '-p', 'results', # prompt text
            '-sep', '\\0',
            '-i', # case-insensitive
            '-markup-rows', # render markup, instead of just displaying it raw
        ]
            + (['-eh', '2'] if search_was_fulltext else []),
        text=True, input=entries_onestring, stdout=PIPE
    )
    if ran.returncode != 0: raise RuntimeError('Rofi (result selection dialog) returned non-zero')

    # process the selection
    selected_idx = int(ran.stdout.split(' ')[0]) # get the index that was prepended to the entry
    return entries[selected_idx]

if __name__ == '__main__':
    args = parse_args() # command-line args

    # get query from user
    query_str = prompt_search_query(fulltext=args.full_text_search, pdf_only=args.pdf_only).strip()
    if query_str == '': raise RuntimeError('empty input')

    # search
    results = search(query_str, fulltext=args.full_text_search, pdf_only=args.pdf_only)

    # have user select file
    selection: SelectableEntry = prompt_select_result(results, args.full_text_search)

    # open the file and exit
    cmd: list[str]
    if (os.path.splitext(selection.file_url)[1] == '.pdf'):
        pdf_cmd: list[str] | None = get_pdf_command(selection.file_url, selection.page)
        if (pdf_cmd is not None): cmd = pdf_cmd
        else: cmd = get_generic_file_command(selection.file_url)
    else:
        cmd = get_generic_file_command(selection.file_url)

    print('Running command: ', end='')
    print(cmd)

    sp.Popen(cmd)

    exit(0) # Popen uses fork() (on Unix systems), so exiting now won't kill the child process