misc.py

import re
import sys
from contextlib import contextmanager
from tempfile import TemporaryDirectory
from typing import *
from urllib.request import urlretrieve

import os

__all__ = [
    'abspath_relative_to_file',
    'fake_tqdm',
    'ensure_parent_exists',
    'as_local_file',
]


def abspath_relative_to_file(path, file_path):
    return os.path.join(
        os.path.split(os.path.abspath(file_path))[0],
        path
    )


def fake_tqdm(data, *args, **kwargs):
    yield from data


def ensure_parent_exists(path):
    if path is not None:
        path = os.path.abspath(path)
        parent_dir = os.path.split(path)[0]
        if not os.path.isdir(parent_dir):
            os.makedirs(parent_dir, exist_ok=True)
        return path


@contextmanager
def as_local_file(uri: str) -> ContextManager[str]:
    if re.match(r'^https?://', uri):
        m = re.match(r'^(https?://[^/]+)/([a-z0-9]{24})/(.*)?$', uri)
        if m:
            uri = f'{m.group(1)}/v1/_getfile/{m.group(2)}'
            if m.group(3):
                uri += f'/{m.group(3)}'
        with TemporaryDirectory() as temp_dir:
            filename = os.path.join(temp_dir, uri.rstrip('/').rsplit('/', 1)[-1])
            print(f'Download: {uri}', file=sys.stderr)
            urlretrieve(uri, filename=filename)
            yield filename
    elif uri.startswith('file://'):
        yield uri[7:]
    else:
        yield uri