marxbook/marxbook/cli.py

#!/usr/bin/env python3
from argparse import ArgumentParser
import os
import subprocess
import sys
import tempfile

import marxbook
import marxbook.extract


def get_parser():
    p = ArgumentParser()
    p.add_argument("--store-dir")
    p.add_argument("--editor", default="sensible-editor")
    p.add_argument("--batch", default=False, action="store_true")
    p.set_defaults(func=None)
    subcommands = p.add_subparsers(help="Sub-commands")
    list_p = subcommands.add_parser("list")
    list_p.add_argument("--format", default="{Path}\t{Tag}\t{Title}\t{Url}", dest='fmt')
    list_p.add_argument("folder", nargs="?", default="")
    list_p.set_defaults(func=main_list)

    add_p = subcommands.add_parser("add")
    add_p.add_argument("--folder", default="")
    add_p.add_argument("--tag", help="Comma-separated list of tags", default="")
    add_p.add_argument("--title", help="If omitted, auto-fetch")
    add_p.add_argument("--description", help="If omitted, auto-fetch")
    add_p.add_argument("url", nargs="?")
    add_p.set_defaults(func=main_add)

    return p


def main():
    p = get_parser()
    args = p.parse_args()

    if args.func is None:
        print("Must specify a subcommand", file=sys.stderr)
        return 2

    store = marxbook.Store(args.store_dir)
    args.func(store, args)


def main_list(store, args):
    for mark in store.folder(args.folder):
        markdata = {"Folder": os.path.dirname(mark['Path'])}
        markdata.update(mark)
        markdata['Tag'] = ','.join(mark['Tag'])
        print(args.fmt.replace(r'\t', '\t').format(**markdata))


def edit_before_add(data: dict, args) -> dict:
    ser = marxbook.Serializer()
    fd, fpath = tempfile.mkstemp()
    buf = os.fdopen(fd, "w")
    buf.write(ser.encode(data))
    buf.close()
    proc = subprocess.Popen([args.editor, fpath])
    proc.communicate()

    with open(fpath) as buf:
        read_data = ser.decode(buf.read())
    os.unlink(fpath)
    data = {}
    for key in read_data:
        data[key.lower()] = read_data[key]
    return data


def main_add(store, args):
    store = store.folder(args.folder)
    batch = args.batch
    if args.url is not None:
        urls = [args.url]
    else:
        batch = True
        urls = []
        for line in sys.stdin.readlines():
            urls.append(line.strip())

    for url in urls:
        data = dict(title=args.title, description=args.description, url=url)
        data['tag'] = [t.strip() for t in args.tag.split(",")]
        if args.title is None or args.description is None:
            _title, _description, _keys, mime, bad = marxbook.extract.network_handler(url)
            if not args.title:
                data["title"] = _title
            if not args.description:
                data["description"] = _description
        if not batch:
            data = edit_before_add(data, args)
        store.add(**data)
    print(urls)


if __name__ == "__main__":
    ret = main()
    if type(ret) is int:
        sys.exit(ret)