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("--edit", dest="edit", action="store_true", default=None)
    add_p.add_argument("--no-edit", dest="edit", action="store_false", default=None)
    add_p.add_argument("url", nargs="?")
    add_p.set_defaults(func=main_add)

    rm_p = subcommands.add_parser("rm")
    rm_p.add_argument("--folder", default="")
    rm_p.add_argument("path")
    rm_p.set_defaults(func=main_rm)

    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))


class Edit:
    def __init__(self, args):
        self.args = args

    def edit_before_add(self, data: dict) -> dict:
        ser = marxbook.Serializer()
        fd, fpath = tempfile.mkstemp()
        buf = os.fdopen(fd, "w")
        buf.write(ser.encode(data))
        buf.close()
        proc = subprocess.Popen([self.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
    editor = Edit(args)
    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
        # shall we edit?
        edit_needed = not batch
        if args.edit is not None:
            edit_needed = args.edit
        if edit_needed:
            data = editor.edit_before_add(data)
        store.add(**data)
    print(urls)


def main_rm(store, args):
    store = store.folder(args.folder)
    dirpath, name = os.path.split(args.path)
    store = store.folder(dirpath)
    store.find(name).unlink()


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