getmanga: Added ability to save node names.

[automanga.git] / getmanga

#!/usr/bin/python3

import sys, os, getopt, time, random
import manga.lib, manga.profile
from PIL import Image

verbose = 0
wait = 10
getnames = False

def msg(vl, msg, *args):
    if verbose >= vl:
        sys.stderr.write("getmanga: " + (msg % args) + "\n")

def getprop(nm, default=None):
    if mprof and "dl-" + nm in mprof.props:
        return mprof.props["dl-" + nm]
    if nm in props:
        return props[nm]
    return default

def digits(num):
    n, i = 10, 1
    while True:
        if num < n:
            return i
        n, i = n * 10, i + 1

def autoname(page):
    ret = ""
    inames = []
    for t, i in page.stack:
        if ret:
            inames.append((t, ret))
        if ret:
            ret += "-"
        ret += "%0*i" % (digits(len(t) + 1), i + 1)
    return ret, inames

def expand(pattern, page):
    ret = ""
    si = 0
    fp = 0
    inames = []
    stack = list(zip([t for t, i in page.stack], [t for t, i in page.stack[1:]] + [page], [i for t, i in page.stack]))
    while True:
        p = pattern.find('%', fp)
        if p < 0:
            if si < len(stack):
                sys.stderr.write("getmanga: pattern %s did not match page %s\n" %
                                 (pattern, "/".join(c.name for t, c, i in stack)))
                sys.exit(1)
            return ret + pattern[fp:], inames
        ret += pattern[fp:p]
        m = pattern[p + 1:p + 2]
        fp = p + 2
        if m == "%":
            ret += "%"
        else:
            if si >= len(stack):
                sys.stderr.write("getmanga: pattern %s did not match page %s\n" %
                                 (pattern, "/".join(c.name for t, c, i in stack)))
                sys.exit(1)
            t, ct, ti = stack[si]
            si += 1
            if m == "i":
                ret += "%0*i" % (digits(len(t) + 1), ti + 1)
            elif m == "n":
                ret += ct.name
            elif m == "d":
                ret += ct.id
            else:
                sys.stderr.write("getmanga: %s: unknown specifier `%s'\n" % (pattern, m))
                sys.exit(1)
            if si < len(stack):
                inames.append((ct, ret))

checkednames = None
def checknames(tdir, names):
    global checkednames
    nmpath = os.path.join(tdir, "names")
    if checkednames is None:
        checkednames = {}
        if os.path.exists(nmpath):
            with open(nmpath) as fp:
                for line in fp:
                    line = line.strip()
                    p = line.find(' ')
                    if p < 0: continue
                    checkednames[line[:p]] = line[p + 1:]
    for t, prefix in names:
        if not prefix: continue
        if ' ' not in prefix and prefix not in checkednames:
            with manga.profile.txfile(nmpath, "w") as fp:
                if '\n' not in t.name:
                    checkednames[prefix] = t.name
                    msg(1, "adding name %s for %s" % (t.name, prefix))
                else:
                    checkednames[prefix] = ""
                    sys.stderr.write("getmanga: warning: node names contains newlines: %r\n" % (t.name,))
                for prefix, name in checkednames.items():
                    if not name: continue
                    fp.write("%s %s\n" % (prefix, name))

def download(mng, tdir, pattern):
    exts = ["", ".jpg", ".jpeg", ".png", ".gif"]
    fmts = {"PNG": "png", "JPEG": "jpeg", "GIF": "gif"}
    for page in manga.lib.cursor(mng):
        if pattern is None:
            nm, inames = autoname(page)
        else:
            nm, inames = expand(pattern, page)
        if getnames:
            checknames(tdir, inames)
        path = os.path.join(tdir, nm)
        if any(os.path.exists(path + ext) for ext in exts):
            msg(2, "%s exists, skipping", nm)
            continue
        msg(1, "getting %s...", nm)
        with page.open() as fp:
            with open(path, "wb") as out:
                done = False
                try:
                    while True:
                        data = fp.read(65536)
                        if data == b"":
                            done = True
                            break
                        out.write(data)
                finally:
                    if not done:
                        os.unlink(path)
            try:
                img = Image.open(path)
            except OSError:
                fmt = None
            else:
                fmt = img.format
            if fmt not in fmts:
                sys.stderr.write("getmanga: warning: could not determine file format of %s, leaving as is\n" % nm)
            else:
                os.rename(path, path + "." + fmts[fmt])
                msg(3, "%s -> %s", nm, nm + "." + fmts[fmt])
        cwait = abs(random.gauss(0, 1) * wait)
        msg(2, "waiting %.1f s...", cwait)
        time.sleep(cwait)

def usage(out):
    out.write("usage: getmanga [-hvn] [-w WAIT] [-p PROFILE] [-P PATTERN] DIRECTORY [LIBRARY ID]\n")
    out.write("\tpattern templates:\n")
    out.write("\t  %i\tSequence number\n")
    out.write("\t  %n\tName\n")
    out.write("\t  %d\tID\n")

def main():
    global verbose, wait, mprof, props, getnames

    opts, args = getopt.getopt(sys.argv[1:], "hvnp:w:P:")
    profnm = None
    pattern = None
    for o, a in opts:
        if o == "-h":
            usage(sys.stdout)
            sys.exit(0)
        elif o == "-p":
            profnm = a
        elif o == "-v":
            verbose += 1
        elif o == "-n":
            getnames = True
        elif o == "-w":
            wait = int(a)
        elif o == "-P":
            pattern = a
    if len(args) < 1:
        usage(sys.stderr)
        sys.exit(1)
    tdir = args[0]

    if not os.path.isdir(tdir):
        sys.stderr.write("getmanga: %s: not a directory\n" % (tdir))
        sys.exit(1)

    pfile = os.path.join(tdir, ".props")
    props = {}
    if os.path.exists(pfile):
        with open(pfile, "r") as fp:
            for words in manga.profile.splitlines(fp):
                if words[0] == "set" and len(words) > 2:
                    props[words[1]] = words[2]
                elif words[0] == "lset" and len(words) > 1:
                    props[words[1]] = words[2:]

    if profnm is None:
        profile = manga.profile.profile.last()
    elif profnm == "":
        profile = None
    else:
        profile = manga.profile.profile.byname(profnm)

    if len(args) == 2:
        usage(sys.stderr)
        sys.exit(1)
    elif len(args) > 2:
        libnm, mid = args[1:3]
    elif isinstance(props.get("manga"), list):
        libnm, mid = props["manga"]
    else:
        sys.stderr.write("getmanga: %s: id is neither saved nor given\n" % (tdir))
        sys.exit(1)
    try:
        lib = manga.lib.findlib(libnm)
    except ImportError:
        sys.stderr.write("getmanga: no such library: %s\n" % (libnm))
        sys.exit(1)
    try:
        mng = lib.byid(mid)
    except KeyError:
        sys.stderr.write("getmanga: no such manga: %s\n" % (mid))
        sys.exit(1)
    if profile is not None:
        mprof = profile.getmanga(libnm, mng.id)
    else:
        mprof = None

    download(mng, tdir, pattern or getprop("pattern"))

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass