Gitlab Community Edition Instance

get.py 3.78 KB
Newer Older
Marcel Hellkamp's avatar
Marcel Hellkamp committed
1
"""
2
Download a single file from an archive.
Marcel Hellkamp's avatar
Marcel Hellkamp committed
3
4
5
6
"""
import os
import sys

7
from pycdstar3.cli import CliError
8
from pycdstar3.cli._utils import hbytes
Marcel Hellkamp's avatar
Marcel Hellkamp committed
9
10


Marcel Hellkamp's avatar
Marcel Hellkamp committed
11
def register(subparsers):
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
    parser = subparsers.add_parser(
        "get", help=__doc__.strip().splitlines()[0], description=__doc__
    )
    parser.add_argument(
        "-%",
        "-p",
        "--progress",
        action="store_true",
        help="Show progress bar for large files or slow downloads",
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="If a <DST>.part file exists, try to resume an"
        " interrupted download. Also, keep the *.part file"
        " on any errors.",
    )
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="Overwrite local files or print binary data to a terminal.",
    )
35
36
    parser.add_argument("ARCHIVE", help="Archive ID")
    parser.add_argument("FILE", help="File Name")
37
38
39
    parser.add_argument(
        "DST",
        nargs="?",
Marcel Hellkamp's avatar
Marcel Hellkamp committed
40
41
        help="Destination filename or directory."
        " No parameter or '-' prints to standard output. (default: '-')",
42
    )
Marcel Hellkamp's avatar
Marcel Hellkamp committed
43
44
45
    parser.set_defaults(main=get)


46
def get(ctx, args):  # noqa: C901
47
48
49
50
51
    client = ctx.client
    vault = ctx.vault
    archive = args.ARCHIVE
    file = args.FILE

Marcel Hellkamp's avatar
Marcel Hellkamp committed
52
53
    resume = args.resume
    progress = args.progress
54
    dst = args.DST or "-"
Marcel Hellkamp's avatar
Marcel Hellkamp committed
55
56
57
58
59
60
61
62
    force = args.force

    if not file:
        raise CliError("The <SRC> parameter must reference a file, not an archive")

    if dst.endswith("/") or os.path.isdir(dst):
        dst = os.path.join(dst, os.path.basename(file))

63
    ispipe = dst == "-"
Marcel Hellkamp's avatar
Marcel Hellkamp committed
64
65
66
67
68
69
70
71
72
    partfile = dst + ".part" if not ispipe else None
    offset = 0

    if not ispipe and not force and os.path.exists(dst):
        raise CliError("File exists: " + dst)

    if partfile and os.path.exists(partfile):
        if resume:
            offset = os.path.getsize(partfile)
73
            ctx.print("Resuming download at: {}", hbytes(offset))
Marcel Hellkamp's avatar
Marcel Hellkamp committed
74
        else:
75
76
77
            raise CliError(
                "Found partial download, but --resume is not set: " + partfile
            )
Marcel Hellkamp's avatar
Marcel Hellkamp committed
78
79
80
81

    if ispipe:
        out = sys.stdout.buffer
    else:
82
        out = open(partfile, "ab")
Marcel Hellkamp's avatar
Marcel Hellkamp committed
83
84
85
86
    write = out.write
    close = out.close

    try:
87
        dl = client.get_file(vault, archive, file, offset=offset)
Marcel Hellkamp's avatar
Marcel Hellkamp committed
88

89
        if ispipe and out.isatty() and not dl.type.startswith("text/") and not force:
90
            raise CliError(
mhellka's avatar
mhellka committed
91
                "Not printing binary data ({}) to a terminal."
Marcel Hellkamp's avatar
Marcel Hellkamp committed
92
                " (--force not set)".format(dl.type)
93
            )
Marcel Hellkamp's avatar
Marcel Hellkamp committed
94

95
        if progress and not ctx.print.quiet:
Marcel Hellkamp's avatar
Marcel Hellkamp committed
96
            from tqdm import tqdm
97
98
99
100
101
102
103
104
105
106

            pbar = tqdm(
                total=dl.size + offset,
                initial=offset,
                unit="b",
                unit_scale=True,
                unit_divisor=1024,
                dynamic_ncols=True,
                file=ctx.print.file,
            )
107

Marcel Hellkamp's avatar
Marcel Hellkamp committed
108
109
110
            def write(chunk):
                pbar.update(len(chunk))
                out.write(chunk)
111

Marcel Hellkamp's avatar
Marcel Hellkamp committed
112
113
114
115
            def close():
                pbar.close()
                out.close()

116
        ctx.print.v("Downloading {} ({})", file, hbytes(dl.size + offset))
Marcel Hellkamp's avatar
Marcel Hellkamp committed
117

118
        for chunk in dl.iter_content():
Marcel Hellkamp's avatar
Marcel Hellkamp committed
119
120
121
122
123
124
125
126
127
            write(chunk)

        close()
        if not ispipe:
            if force:
                os.replace(partfile, dst)
            else:
                os.rename(partfile, dst)

128
        ctx.print.v("Done!", file, vault, archive, hbytes(dl.size + offset))
Marcel Hellkamp's avatar
Marcel Hellkamp committed
129

130
    except (KeyboardInterrupt, Exception):
Marcel Hellkamp's avatar
Marcel Hellkamp committed
131
132
133
134
135
136
        if close:
            close()
        if partfile and os.path.exists(partfile) and not resume:
            try:
                os.remove(partfile)
            except OSError:
137
                ctx.print("Failed to delete partial download: {}", partfile)
Marcel Hellkamp's avatar
Marcel Hellkamp committed
138
139
                raise
        raise