"""Utility functions used in the fastai library"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_xtras.ipynb.

# %% ../nbs/03_xtras.ipynb #ecd054a6
from __future__ import annotations

# %% auto #0
__all__ = ['spark_chars', 'UNSET', 'walk', 'exttypes', 'globtastic', 'pglob', 'maybe_open', 'mkdir', 'image_size', 'img_bytes',
           'detect_mime', 'bunzip', 'loads', 'loads_multi', 'dumps', 'untar_dir', 'repo_details', 'shell', 'ssh',
           'rsync_multi', 'run', 'open_file', 'save_pickle', 'load_pickle', 'parse_env', 'expand_wildcards',
           'atomic_save', 'dict2obj', 'obj2dict', 'repr_dict', 'is_listy', 'mapped', 'IterLen', 'ReindexCollection',
           'SaveReturn', 'trim_wraps', 'save_iter', 'asave_iter', 'frontmatter', 'clean_cli_output', 'unqid',
           'rtoken_hex', 'friendly_name', 'n_friendly_names', 'exec_eval', 'get_source_link', 'sparkline',
           'modify_exception', 'round_multiple', 'set_num_threads', 'join_path_file', 'autostart', 'EventTimer',
           'stringfmt_names', 'PartialFormatter', 'partial_format', 'truncstr', 'utc2local', 'local2utc', 'trace',
           'modified_env', 'ContextManagers', 'shufflish', 'console_help', 'hl_md', 'type2str', 'dataclass_src',
           'Unset', 'nullable_dc', 'make_nullable', 'flexiclass', 'asdict', 'vars_pub', 'is_typeddict', 'is_namedtuple',
           'CachedIter', 'CachedAwaitable', 'reawaitable', 'is_async_callable', 'maybe_await', 'noopa', 'flexicache',
           'time_policy', 'mtime_policy', 'timed_cache']

# %% ../nbs/03_xtras.ipynb #3401d507
from .imports import *
from .foundation import *
from .basics import *
from .ansi import strip_ansi

from importlib import import_module
from functools import wraps
import string,time,dataclasses
from enum import Enum
from contextlib import contextmanager,ExitStack
from datetime import datetime, timezone
from time import sleep,time,perf_counter
from os.path import getmtime
from dataclasses import dataclass, field, fields, is_dataclass, MISSING, make_dataclass

# %% ../nbs/03_xtras.ipynb #f64c2b09
def walk(
    path:Path|str, # path to start searching
    symlinks:bool=True, # follow symlinks?
    keep_file:callable=ret_true, # function that returns True for wanted files
    keep_folder:callable=ret_true, # function that returns True for folders to enter
    skip_folder:callable=ret_false, # function that returns True for folders to skip
    func:callable=os.path.join, # function to apply to each matched file
    ret_folders:bool=False, # return folders, not just files
    sort:bool=True # sort files by name within each folder
):
    "Generator version of `os.walk`, using functions to filter files and folders"
    from copy import copy
    for root,dirs,files in os.walk(path, followlinks=symlinks):
        if keep_folder(root,''):
            if ret_folders: yield func(root, '')
            if sort: files = sorted(files)
            yield from (func(root, name) for name in files if keep_file(root,name))
        for name in copy(dirs):
            if skip_folder(root,name): dirs.remove(name)

# %% ../nbs/03_xtras.ipynb #12a01ff4
_exttypes = dict(
    py=['ipynb', 'py'],
    js=['js', 'jsx', 'mjs', 'cjs'],
    java=['java', 'kt', 'scala', 'gradle'],
    c=['c', 'h'],
    cpp=['cpp', 'hpp', 'cc', 'cxx'],
    rb=['rb', 'rake'],
    r=['r', 'R'],
    ex=['ex', 'exs', 'erl', 'hrl'],
    sh=['sh', 'bash'],
    web=['html', 'css'],
    doc=['md', 'rst'],
    cfg=['json', 'yaml', 'yml', 'toml', 'xml'],
)

def exttypes(types):
    """Get exts for comma-separated or list `typ`; if not found in list, return list with just `types`.
    Supported: py, js, java, c, cpp, rb, r, ex, sh, web, doc, cfg"""
    typs = L(types).flatmap(Self.split(','))
    return typs.flatmap(lambda x: _exttypes.get(x, [x]))

# %% ../nbs/03_xtras.ipynb #4fef7fef
def globtastic(
    path:Path|str='.', # path to start searching
    recursive:bool=True, # search subfolders
    symlinks:bool=True, # follow symlinks?
    file_glob:str=None, # Only include files matching glob
    file_re:str=None, # Only include files matching regex
    folder_re:str=None, # Only enter folders matching regex
    skip_file_glob:str=None, # Skip files matching glob
    skip_file_re:str=None, # Skip files matching regex
    skip_folder_re:str=None, # Skip folders matching regex,
    func:callable=os.path.join, # function to apply to each matched file
    ret_folders:bool=False, # return folders, not just files
    sort:bool=True, # sort files by name within each folder
    types:str|list=None,  # list or comma-separated str of ext types from: py, js, java, c, cpp, rb, r, ex, sh, web, doc, cfg
    exts:str|list=None, # list or comma-separated str of exts to include
)->L: # Paths to matched files
    "A more powerful `glob`, including regex matches, symlink handling, and skip parameters"
    from fnmatch import fnmatch
    path = Path(path)
    if path.is_file(): return L([path])
    if not recursive: skip_folder_re='.'
    exts = L(exts).flatmap(Self.split(','))
    exts += exttypes(types)
    if exts:
        exts = [e if e.startswith('.') else f'.{e}' for e in listify(exts)]
        file_re = f"({'|'.join(re.escape(e) for e in exts)})$"
    file_re,folder_re = compile_re(file_re),compile_re(folder_re)
    skip_file_re,skip_folder_re = compile_re(skip_file_re),compile_re(skip_folder_re)
    def _keep_file(root, name):
        return (not file_glob or fnmatch(name, file_glob)) and (
                not file_re or file_re.search(name)) and (
                not skip_file_glob or not fnmatch(name, skip_file_glob)) and (
                not skip_file_re or not skip_file_re.search(name))
    def _keep_folder(root, name): return not folder_re or folder_re.search(os.path.join(root,name))
    def _skip_folder(root, name): return skip_folder_re and skip_folder_re.search(name)
    return L(walk(path, symlinks=symlinks, keep_file=_keep_file, keep_folder=_keep_folder, skip_folder=_skip_folder,
                  func=func, ret_folders=ret_folders, sort=sort))

# %% ../nbs/03_xtras.ipynb #827b4914
@fdelegates(globtastic)
def pglob(
    path:Path|str='.', # path to start searching
    func:callable=Path, # function to apply to each matched file
    **kwargs
)->L: # Paths to matched files
    "Shortcut for `globtastic(..., call=Path)`"
    return globtastic(path, func=func, **kwargs)

# %% ../nbs/03_xtras.ipynb #d514da68
@contextmanager
def maybe_open(f, mode='r', **kwargs):
    "Context manager: open `f` if it is a path (and close on exit)"
    if isinstance(f, (str,os.PathLike)):
        with open(f, mode, **kwargs) as f: yield f
    else: yield f

# %% ../nbs/03_xtras.ipynb #63d86576
def mkdir(path, exist_ok=False, parents=False, overwrite=False, **kwargs):
    "Creates and returns a directory defined by `path`, optionally removing previous existing directory if `overwrite` is `True`"
    import shutil
    path = Path(path)
    if path.exists() and overwrite: shutil.rmtree(path)
    path.mkdir(exist_ok=exist_ok, parents=parents, **kwargs)
    return path

# %% ../nbs/03_xtras.ipynb #b33313ff
def image_size(fn):
    "Tuple of (w,h) for png, gif, or jpg; `None` otherwise"
    from fastcore import imghdr
    import struct
    def _jpg_size(f):
        size,ftype = 2,0
        while not 0xc0 <= ftype <= 0xcf:
            f.seek(size, 1)
            byte = f.read(1)
            while ord(byte) == 0xff: byte = f.read(1)
            ftype = ord(byte)
            size = struct.unpack('>H', f.read(2))[0] - 2
        f.seek(1, 1)  # `precision'
        h,w = struct.unpack('>HH', f.read(4))
        return w,h

    def _gif_size(f): return struct.unpack('<HH', head[6:10])

    def _png_size(f):
        assert struct.unpack('>i', head[4:8])[0]==0x0d0a1a0a
        return struct.unpack('>ii', head[16:24])
    d = dict(png=_png_size, gif=_gif_size, jpeg=_jpg_size)
    with maybe_open(fn, 'rb') as f: return d[imghdr.what(f)](f)

# %% ../nbs/03_xtras.ipynb #d3aa6edf
def img_bytes(img, fmt='PNG'):
    # Convert PIL `img` to bytes in format `fmt`
    from io import BytesIO
    buf=BytesIO()
    img.save(buf, format=fmt)
    return buf.getvalue()

# %% ../nbs/03_xtras.ipynb #e5801198
_sigs = {
    (b'%PDF', 0): 'application/pdf',
    (b'RIFF', 0): lambda d: 'audio/wav' if d[8:12]==b'WAVE' else 'video/avi' if d[8:12]==b'AVI ' else None,
    (b'ID3', 0): 'audio/mp3',
    (b'\xff\xfb', 0): 'audio/mp3',
    (b'\xff\xf3', 0): 'audio/mp3',
    (b'FORM', 0): lambda d: 'audio/aiff' if d[8:12]==b'AIFF' else None,
    (b'OggS', 0): 'audio/ogg',
    (b'fLaC', 0): 'audio/flac',
    (b'ftyp', 4): lambda d: 'video/3gpp' if d[8:11]==b'3gp' else 'video/mp4',
    (b'\x1a\x45\xdf', 0): 'video/webm',
    (b'FLV', 0): 'video/x-flv',
    (b'\x30\x26\xb2\x75', 0): 'video/wmv',
    (b'\x00\x00\x01\xb3', 0): 'video/mpeg',
}

def detect_mime(data):
    "Get the MIME type for bytes `data`, covering common PDF, audio, video, and image types"
    if not isinstance(data,bytes): return None
    import mimetypes
    from fastcore import imghdr
    for (sig,pos),mime in _sigs.items():
        if data[pos:pos+len(sig)]==sig: return mime(data) if callable(mime) else mime
    return mimetypes.types_map.get(f'.{imghdr.what(None, h=data)}')

# %% ../nbs/03_xtras.ipynb #f856746a
def bunzip(fn):
    "bunzip `fn`, raising exception if output already exists"
    fn = Path(fn)
    assert fn.exists(), f"{fn} doesn't exist"
    out_fn = fn.with_suffix('')
    assert not out_fn.exists(), f"{out_fn} already exists"
    import bz2
    with bz2.BZ2File(fn, 'rb') as src, out_fn.open('wb') as dst:
        for d in iter(lambda: src.read(1024*1024), b''): dst.write(d)

# %% ../nbs/03_xtras.ipynb #0a78c637
def loads(s, **kw):
    "Same as `json.loads`, but handles `None`"
    if not s: return {}
    try: import ujson as json
    except ModuleNotFoundError: import json
    return json.loads(s, **kw)

# %% ../nbs/03_xtras.ipynb #131170f7
def loads_multi(s:str):
    "Generator of >=0 decoded json dicts, possibly with non-json ignored text at start and end"
    import json
    _dec = json.JSONDecoder()
    while s.find('{')>=0:
        s = s[s.find('{'):]
        obj,pos = _dec.raw_decode(s)
        if not pos: raise ValueError(f'no JSON object found at {pos}')
        yield obj
        s = s[pos:]

# %% ../nbs/03_xtras.ipynb #ed688ac4
def dumps(obj, **kw):
    "Same as `json.dumps`, but uses `ujson` if available"
    try: import ujson as json
    except ModuleNotFoundError: import json
    else: kw['escape_forward_slashes']=False
    return json.dumps(obj, **kw)

# %% ../nbs/03_xtras.ipynb #ac84300c
def _unpack(fname, out):
    import shutil
    shutil.unpack_archive(str(fname), str(out))
    ls = out.ls()
    return ls[0] if len(ls) == 1 else out

# %% ../nbs/03_xtras.ipynb #6f1e90e9
def untar_dir(fname, dest, rename=False, overwrite=False):
    "untar `file` into `dest`, creating a directory if the root contains more than one item"
    import tempfile,shutil
    dest = Path(dest)
    with tempfile.TemporaryDirectory() as d:
        out = Path(d)/remove_suffix(Path(fname).stem, '.tar')
        out.mkdir()
        if rename: dest = dest/out.name
        else:
            src = _unpack(fname, out)
            dest = dest/src.name
        if dest.exists():
            if overwrite: shutil.rmtree(dest) if dest.is_dir() else dest.unlink()
            else: return dest
        if rename: src = _unpack(fname, out)
        shutil.move(str(src), dest)
        return dest

# %% ../nbs/03_xtras.ipynb #ab91ee87
def repo_details(url):
    "Tuple of `owner,name` from ssh or https git repo `url`"
    res = remove_suffix(url.strip(), '.git')
    res = res.split(':')[-1]
    return res.split('/')[-2:]

# %% ../nbs/03_xtras.ipynb #ba70e535
def shell(*args, **kwargs):
    "Shortcut for `subprocess.run(shell=True)`"
    import subprocess
    return subprocess.run(*args, shell=True, **kwargs)

# %% ../nbs/03_xtras.ipynb #f62134a5
def ssh(host, args='', user='ubuntu', sock=None):
    "Run SSH command with given arguments"
    sock_opts = f'-S {sock}' if sock else ''
    return shell(f'ssh {sock_opts} {args} {user}@{host}')

# %% ../nbs/03_xtras.ipynb #cb24780e
def rsync_multi(ip, files, user='ubuntu', persist='5m'):
    "Transfer multiple files with rename using persistent SSH connection"
    sock = f'/tmp/ssh-{ip}-{user}'
    ssh(ip, f'-o ControlMaster=auto -o ControlPersist={persist} -N -f', user, sock)
    for src,dst in files: shell(f'rsync -az -e "ssh -S {sock}" {src} {user}@{ip}:{dst}')

# %% ../nbs/03_xtras.ipynb #dda52646
def run(cmd, *rest, same_in_win=False, ignore_ex=False, as_bytes=False, stderr=True):
    "Pass `cmd` (splitting with `shlex` if string) to `subprocess.run`; return `stdout`; raise `IOError` if fails"
    # Even the command is same on Windows, we have to add `cmd /c `"
    import subprocess
    if rest:
        if sys.platform == 'win32' and same_in_win:
            cmd = ('cmd', '/c', cmd, *rest)
        else:
            cmd = (cmd,)+rest
    elif isinstance(cmd, str):
        if sys.platform == 'win32' and same_in_win: cmd = 'cmd /c ' + cmd
        import shlex
        cmd = shlex.split(cmd)
    elif isinstance(cmd, list):
        if sys.platform == 'win32' and same_in_win: cmd = ['cmd', '/c'] + cmd
    res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout = res.stdout
    if stderr and res.stderr: stdout += b' ;; ' + res.stderr
    if not as_bytes: stdout = stdout.decode().strip()
    if ignore_ex: return (res.returncode, stdout)
    if res.returncode: raise IOError((res.stdout + b' ;; ' + res.stderr).decode().strip())
    return stdout

# %% ../nbs/03_xtras.ipynb #5097b830
def open_file(fn, mode='r', **kwargs):
    "Open a file, with optional compression if gz or bz2 suffix"
    if isinstance(fn, io.IOBase): return fn
    import bz2,gzip,zipfile
    fn = Path(fn)
    if   fn.suffix=='.bz2': return bz2.BZ2File(fn, mode, **kwargs)
    elif fn.suffix=='.gz' : return gzip.GzipFile(fn, mode, **kwargs)
    elif fn.suffix=='.zip': return zipfile.ZipFile(fn, mode, **kwargs)
    else: return open(fn,mode, **kwargs)

# %% ../nbs/03_xtras.ipynb #a9fc7a40
def save_pickle(fn, o):
    "Save a pickle file, to a file name or opened file"
    import pickle
    with open_file(fn, 'wb') as f: pickle.dump(o, f)

# %% ../nbs/03_xtras.ipynb #5e9bc30e
def load_pickle(fn):
    "Load a pickle file from a file name or opened file"
    import pickle
    with open_file(fn, 'rb') as f: return pickle.load(f)

# %% ../nbs/03_xtras.ipynb #02029d02
def parse_env(s:str=None, fn:Union[str,Path]=None) -> dict:
    "Parse a shell-style environment string or file"
    assert bool(s)^bool(fn), "Must pass exactly one of `s` or `fn`"
    if fn: s = Path(fn).read_text()
    def _f(line):
        m = re.match(r'^\s*(?:export\s+)?(\w+)\s*=\s*(["\']?)(.*?)(\2)\s*(?:#.*)?$', line).groups()
        return m[0], m[2]

    return dict(_f(o.strip()) for o in s.splitlines() if o.strip() and not re.match(r'\s*#', o))

# %% ../nbs/03_xtras.ipynb #d030f939
def expand_wildcards(code):
    "Expand all wildcard imports in the given code string."
    import ast,importlib
    tree = ast.parse(code)

    def _replace_node(code, old_node, new_node):
        "Replace `old_node` in the source `code` with `new_node`."
        lines = code.splitlines()
        lnum = old_node.lineno
        indent = ' ' * (len(lines[lnum-1]) - len(lines[lnum-1].lstrip()))
        new_lines = [indent+line for line in ast.unparse(new_node).splitlines()]
        lines[lnum-1 : old_node.end_lineno] = new_lines
        return '\n'.join(lines)

    def _expand_import(node, mod, existing):
        "Create expanded import `node` in `tree` from wildcard import of `mod`."
        mod_all = getattr(mod, '__all__', None)
        available_names = set(mod_all) if mod_all is not None else set(dir(mod))
        used_names = {n.id for n in ast.walk(tree) if isinstance(n, ast.Name) and n.id in available_names} - existing
        if not used_names: return node
        names = [ast.alias(name=name, asname=None) for name in sorted(used_names)]
        return ast.ImportFrom(module=node.module, names=names, level=node.level)

    existing = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom) and node.names[0].name != '*': existing.update(n.name for n in node.names)
        elif isinstance(node, ast.Import): existing.update(n.name.split('.')[0] for n in node.names)
    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom) and any(n.name == '*' for n in node.names):
            new_import = _expand_import(node, importlib.import_module(node.module), existing)
            code = _replace_node(code, node, new_import)
    return code

# %% ../nbs/03_xtras.ipynb #15c2b62c
@patch
def mkdir_perms(self:Path, mode=0o777, parents=False, exist_ok=False, uid=-1, gid=-1):
    "Create directory like Path.mkdir but optionally set uid/gid on newly created dirs"
    if self.exists() and not exist_ok: raise FileExistsError(self)
    to_create = []
    p = self
    while not p.exists():
        to_create.append(p)
        if parents: p = p.parent
        else: break
    for d in reversed(to_create):
        try: d.mkdir(mode=mode)
        except FileExistsError:
            if d==self and not exist_ok: raise
            continue
        if uid>-1 or gid>-1: os.chown(d, uid, gid)

# %% ../nbs/03_xtras.ipynb #832397a3
@contextmanager
def atomic_save(fn, mode='wb', uid=-1, gid=-1, **kwargs):
    "Context manager for writing a file atomically via a temp file that is renamed on close"
    from tempfile import NamedTemporaryFile as ntf
    fn = Path(fn)
    fn.parent.mkdir_perms(parents=True, exist_ok=True, uid=uid, gid=gid)
    with ntf(mode=mode, dir=fn.parent, delete=False, suffix=fn.suffix, **kwargs) as f:
        try: yield f
        except:
            Path(f.name).unlink(missing_ok=True)
            raise
    if uid>-1 or gid>-1: os.chown(f.name, uid, gid)
    Path(f.name).rename(fn)


# %% ../nbs/03_xtras.ipynb #9579358d
def dict2obj(d=None, list_func=L, dict_func=AttrDict, **kwargs):
    "Convert (possibly nested) dicts (or lists of dicts) to `AttrDict`"
    if d is None: d={}
    if isinstance(d, (L,list)): return list_func([dict2obj(v, list_func=list_func, dict_func=dict_func) for v in d])
    if not isinstance(d, dict): return d
    return dict_func(**{k:dict2obj(v, list_func=list_func, dict_func=dict_func) for k,v in (d|kwargs).items()})

# %% ../nbs/03_xtras.ipynb #055a1e07
def obj2dict(d):
    "Convert (possibly nested) AttrDicts (or lists of AttrDicts) to `dict`"
    if isinstance(d, (L,list)): return list(L(d).map(obj2dict))
    if not isinstance(d, dict): return d
    return dict(**{k:obj2dict(v) for k,v in d.items()})

# %% ../nbs/03_xtras.ipynb #36b00c72
def _repr_dict(d, lvl):
    if isinstance(d,dict):
        its = [f"{k}: {_repr_dict(v,lvl+1)}" for k,v in d.items()]
    elif isinstance(d,(list,L)): its = [_repr_dict(o,lvl+1) for o in d]
    else: return str(d)
    return '\n' + '\n'.join([" "*(lvl*2) + "- " + o for o in its])

# %% ../nbs/03_xtras.ipynb #13bb8fba
def repr_dict(d):
    "Print nested dicts and lists, such as returned by `dict2obj`"
    return _repr_dict(d,0).strip()

# %% ../nbs/03_xtras.ipynb #20842e1f
def is_listy(x):
    "`isinstance(x, (tuple,list,L,slice,Generator))`"
    return isinstance(x, (tuple,list,L,slice,Generator))

# %% ../nbs/03_xtras.ipynb #40c2c641
def mapped(f, it):
    "map `f` over `it`, unless it's not listy, in which case return `f(it)`"
    return L(it).map(f) if is_listy(it) else f(it)

# %% ../nbs/03_xtras.ipynb #ab0758e6
@patch
def readlines(self:Path, hint=-1, encoding='utf8'):
    "Read the content of `self`"
    with self.open(encoding=encoding) as f: return f.readlines(hint)

# %% ../nbs/03_xtras.ipynb #5ad0571d
@patch
def read_json(self:Path, encoding=None, errors=None):
    "Same as `read_text` followed by `loads`"
    return loads(self.read_text(encoding=encoding, errors=errors))

# %% ../nbs/03_xtras.ipynb #d6d8d893
@patch
def mk_write(self:Path, data, encoding=None, errors=None, mode=511, uid=-1, gid=-1):
    "Make all parent dirs of `self`, and write `data`"
    self.parent.mkdir(exist_ok=True, parents=True, mode=mode)
    self.write_text(data, encoding=encoding, errors=errors)
    if uid!=-1 or gid!=-1: os.chown(self, uid, gid)

# %% ../nbs/03_xtras.ipynb #233a8732
@patch
def write_json(self:Path, data, encoding=None, errors=None, mode=511, uid=-1, gid=-1, **kw):
    "Same as `dumps`followed by `mk_write`"
    self.mk_write(dumps(data,**kw),encoding,errors,mode,uid,gid)

# %% ../nbs/03_xtras.ipynb #9fc9965f
@patch
def relpath(self:Path, start=None):
    "Same as `os.path.relpath`, but returns a `Path`, and resolves symlinks"
    return Path(os.path.relpath(self.resolve(), Path(start).resolve()))

# %% ../nbs/03_xtras.ipynb #12264861
@patch
def ls(self:Path, n_max=None, file_type=None, file_exts=None):
    "Contents of path as a list"
    import mimetypes
    extns=L(file_exts)
    if file_type: extns += L(k for k,v in mimetypes.types_map.items() if v.startswith(file_type+'/'))
    has_extns = len(extns)==0
    res = (o for o in self.iterdir() if has_extns or o.suffix in extns)
    if n_max is not None: res = itertools.islice(res, n_max)
    return L(res)

# %% ../nbs/03_xtras.ipynb #abe052c6
@patch
def normpath(self:Path):
    "Normalize path, eliminating double slashes, etc."
    return Path(os.path.normpath(self))

# %% ../nbs/03_xtras.ipynb #695f90e0
@patch
def __repr__(self:Path):
    b = getattr(Path, 'BASE_PATH', None)
    if b:
        try: self = self.relative_to(b)
        except: pass
    return f"Path({self.as_posix()!r})"

# %% ../nbs/03_xtras.ipynb #6bd8b653
@patch
def delete(self:Path):
    "Delete a file, symlink, or directory tree"
    if not self.exists(): return
    if self.is_dir():
        import shutil
        shutil.rmtree(self)
    else: self.unlink()

# %% ../nbs/03_xtras.ipynb #f2bb033e
class IterLen:
    "Base class to add iteration to anything supporting `__len__` and `__getitem__`"
    def __iter__(self): return (self[i] for i in range_of(self))

# %% ../nbs/03_xtras.ipynb #0356640b
@docs
class ReindexCollection(GetAttr, IterLen):
    "Reindexes collection `coll` with indices `idxs` and optional LRU cache of size `cache`"
    _default='coll'
    def __init__(self, coll, idxs=None, cache=None, tfm=noop):
        if idxs is None: idxs = L.range(coll)
        store_attr()
        if cache is not None: self._get = functools.lru_cache(maxsize=cache)(self._get)

    def _get(self, i): return self.tfm(self.coll[i])
    def __getitem__(self, i): return self._get(self.idxs[i])
    def __len__(self): return len(self.coll)
    def reindex(self, idxs): self.idxs = idxs
    def shuffle(self):
        import random
        random.shuffle(self.idxs)
    def cache_clear(self): self._get.cache_clear()
    def __getstate__(self): return {'coll': self.coll, 'idxs': self.idxs, 'cache': self.cache, 'tfm': self.tfm}
    def __setstate__(self, s): self.coll,self.idxs,self.cache,self.tfm = s['coll'],s['idxs'],s['cache'],s['tfm']

    _docs = dict(reindex="Replace `self.idxs` with idxs",
                shuffle="Randomly shuffle indices",
                cache_clear="Clear LRU cache")

# %% ../nbs/03_xtras.ipynb #59a2f54d
class SaveReturn:
    "Wrap an iterator such that the generator function's return value is stored in `.value`"
    def __init__(self, its): self.its = its
    def __iter__(self):
        self.value = yield from self.its
        return self.value

# %% ../nbs/03_xtras.ipynb #e89ff9b0
def trim_wraps(f, n=1):
    "Like wraps, but removes the first n parameters from the signature"
    import inspect
    def _(g):
        g = wraps(f)(g)
        sig = inspect.signature(f)
        params = list(sig.parameters.values())[n:]
        g.__signature__ = sig.replace(parameters=params)
        return g
    return _

# %% ../nbs/03_xtras.ipynb #b867efdd
class _save_iter:
    def __init__(self, g, *args, **kw): self.g,self.args,self.kw = g,args,kw
    def __iter__(self): yield from self.g(self, *self.args, **self.kw)
    def __aiter__(self): return self.g(self, *self.args, **self.kw)

def save_iter(g):
    "Decorator that allows a generator function to store values in the returned iterator object"
    @trim_wraps(g)
    def _(*args, **kwargs): return _save_iter(g, *args, **kwargs)
    return _

# %% ../nbs/03_xtras.ipynb #d4017861
def asave_iter(g):
    "Like `save_iter`, but for async iterators"
    @trim_wraps(g)
    def _(*args, **kwargs): return _save_iter(g, *args, **kwargs)
    return _

# %% ../nbs/03_xtras.ipynb #20e906c7
def frontmatter(txt:str)->dict:
    "Dict contained in frontmatter in `txt`, if present"
    import yaml
    if not txt.startswith('---\n'): return {}
    fm,part,_ = txt[4:].partition('\n---\n')
    if not part: return {}
    try: res = yaml.safe_load(fm)
    except yaml.parser.ParserError: return {}
    return res if isinstance(res,dict) else {}

# %% ../nbs/03_xtras.ipynb #45eb5141
def clean_cli_output(txt:str, strip:bool=True):
    "Clean CLI output by handling alternate screen, carriage returns, and ANSI escapes"
    if '\x1b[?1049h' in txt:
        if '\x1b[?1049l' not in txt: return ''
        txt = txt.rsplit('\x1b[?1049l', 1)[-1]
    txt = txt.replace('\r\n', '\n')
    res = '\n'.join(l.rsplit('\r', 1)[-1] for l in txt.split('\n'))
    return strip_ansi(res, term_queries=True) if strip else res

# %% ../nbs/03_xtras.ipynb #35368de5
def unqid(seeded=False):
    "Generate a unique id suitable for use as a Python identifier"
    from base64 import b64encode
    from uuid import uuid4,UUID
    id4 = UUID(int=random.getrandbits(128), version=4) if seeded else uuid4()
    res = b64encode(id4.bytes)
    return '_' + res.decode().rstrip('=').translate(str.maketrans('+/', '_-'))

# %% ../nbs/03_xtras.ipynb #e0e7b294
def rtoken_hex(
    nbytes:int=16 # Number of bytes to generate
)->str: # hex string of length nbytes*2
    "Generate a random hex string using Python's random module."
    import random
    return ''.join(f'{random.randint(0, 255):02x}' for _ in range(nbytes))

# %% ../nbs/03_xtras.ipynb #4e02b1f0
def friendly_name(levels=3, suffix=4):
    "Generate a random human-readable name with customizable word levels and suffix length"
    import random,string
    adjectives = ['admiring', 'adoring', 'amazing', 'awesome', 'beautiful', 'blissful', 'bold', 'brave', 'busy', 'charming', 'clever', 'compassionate', 'confident', 'cool', 'dazzling', 'determined', 'dreamy', 'eager', 'ecstatic', 'elastic', 'elated', 'elegant', 'epic', 'exciting', 'fervent', 'festive', 'flamboyant', 'focused', 'friendly', 'frosty', 'funny', 'gallant', 'gifted', 'goofy', 'gracious', 'great', 'happy', 'hopeful', 'hungry', 'inspiring', 'intelligent', 'interesting', 'jolly', 'jovial', 'keen', 'kind', 'laughing', 'loving', 'lucid', 'magical', 'modest', 'nice', 'nifty', 'nostalgic', 'objective', 'optimistic', 'peaceful', 'pensive', 'practical', 'priceless', 'quirky', 'quizzical', 'relaxed', 'reverent', 'romantic', 'serene', 'sharp', 'silly', 'sleepy', 'stoic', 'sweet', 'tender', 'trusting', 'upbeat', 'vibrant', 'vigilant', 'vigorous', 'wizardly', 'wonderful', 'youthful', 'zealous', 'zen', 'golden', 'silver', 'crimson', 'azure', 'emerald', 'violet', 'amber', 'coral', 'turquoise', 'lavender', 'minty', 'citrus', 'vanilla', 'woody', 'floral', 'fresh', 'gentle', 'sparkling', 'precise', 'curious']
    nouns = ['tiger', 'eagle', 'river', 'mountain', 'forest', 'ocean', 'star', 'moon', 'wind', 'dragon', 'phoenix', 'wolf', 'bear', 'lion', 'shark', 'falcon', 'raven', 'crystal', 'diamond', 'ruby', 'sapphire', 'pearl', 'wave', 'tide', 'cloud', 'rainbow', 'sunset', 'sunrise', 'galaxy', 'comet', 'meteor', 'planet', 'nebula', 'cosmos', 'universe', 'atom', 'photon', 'quantum', 'matrix', 'cipher', 'code', 'signal', 'pulse', 'beam', 'ray', 'spark', 'frost', 'ice', 'snow', 'mist', 'fog', 'dew', 'rain', 'hail', 'helix', 'prism', 'lens', 'mirror', 'echo', 'heart', 'mind', 'dream', 'vision', 'hope', 'wish', 'magic', 'spell', 'charm', 'rune', 'symbol', 'token', 'key', 'door', 'gate', 'bridge', 'tower', 'castle', 'fortress', 'shield', 'dolphin', 'whale', 'penguin', 'butterfly', 'hummingbird', 'deer', 'rabbit', 'fox', 'otter', 'panda', 'koala', 'zebra', 'giraffe', 'elephant', 'valley', 'canyon', 'meadow', 'prairie', 'island', 'lake', 'pond', 'stream', 'waterfall', 'cliff', 'peak', 'hill', 'grove', 'garden', 'sunlight', 'breeze', 'melody', 'sparkle', 'whirlpool', 'windmill', 'carousel', 'spiral', 'glow']
    verbs = ['runs', 'flies', 'jumps', 'builds', 'creates', 'flows', 'shines', 'grows', 'moves', 'works', 'dances', 'sings', 'plays', 'dreams', 'thinks', 'learns', 'teaches', 'helps', 'heals', 'saves', 'protects', 'guards', 'watches', 'sees', 'hears', 'feels', 'knows', 'understands', 'discovers', 'explores', 'searches', 'finds', 'seeks', 'holds', 'carries', 'lifts', 'pushes', 'pulls', 'makes', 'crafts', 'forges', 'shapes', 'forms', 'molds', 'carves', 'joins', 'connects', 'links', 'binds', 'ties', 'opens', 'closes', 'starts', 'stops', 'begins', 'ends', 'finishes', 'completes', 'wins', 'triumphs', 'succeeds', 'achieves', 'accomplishes', 'reaches', 'arrives', 'departs', 'leaves', 'returns', 'comes', 'goes', 'travels', 'journeys', 'walks', 'sprints', 'races', 'speeds', 'rushes', 'hurries', 'waits', 'pauses', 'rests', 'sleeps', 'wakes', 'rises', 'climbs', 'ascends', 'descends', 'swims', 'dives', 'surfs', 'sails', 'paddles', 'hikes', 'treks', 'wanders', 'roams', 'ventures', 'navigates', 'glides', 'soars', 'floats', 'drifts', 'tosses', 'divides', 'shares', 'secures', 'settles', 'places', 'wonders', 'questions']
    adverbs = ['quickly', 'gracefully', 'gently', 'boldly', 'quietly', 'swiftly', 'carefully', 'eagerly', 'smoothly', 'brightly', 'softly', 'steadily', 'cleverly', 'proudly', 'calmly', 'freely', 'wisely', 'kindly', 'firmly', 'lightly', 'deeply', 'clearly', 'warmly', 'coolly', 'sharply', 'slowly', 'rapidly', 'silently', 'loudly', 'perfectly']
    suffix = ''.join(random.choices(string.ascii_lowercase + string.digits, k=suffix))
    words = [random.choice(o) for o in (adjectives,nouns,verbs,adverbs)[:levels]]
    return f"{'-'.join(words)}-{suffix}"

# %% ../nbs/03_xtras.ipynb #998e08a0
def n_friendly_names(levels=3, suffix=4):
    "Number of possible combos for `friendly_names"
    ns = [102,116,110,30]
    return product(ns[:levels])*36**suffix

# %% ../nbs/03_xtras.ipynb #b4d1831f
def exec_eval(code,   # Code to exec/eval
              g=None, # Globals namespace dict
              l=None  # Locals namespace dict
             ):
    "Evaluate `code` in `g` (defaults to `globals()`) and `l` (defaults to `locals()`)"
    import ast, inspect
    frame = inspect.currentframe().f_back
    if g is None: g = frame.f_globals
    if l is None: l = g
    tree = ast.parse(code, mode='exec')
    if tree.body and isinstance(tree.body[-1], ast.Expr):
        *statements, expr = tree.body
        exec_code = compile(ast.Module(statements, []), filename="<string>", mode="exec")
        expr_code = compile(ast.Expression(expr.value), filename="<string>", mode="eval")
        exec(exec_code, g, l)
        return eval(expr_code, g, l)
    else: exec(code, g, l)

# %% ../nbs/03_xtras.ipynb #bdaa78a0
def _unwrapped_type_dispatch_func(x): 
    # use isinstance_str to avoid adding plum-dispatch as dependency to fastcore
    return x.methods[0].implementation if isinstance_str(x,"Function") else x

# %% ../nbs/03_xtras.ipynb #bb4d3fe2
def _is_property(x): return type(x)==property
def _has_property_getter(x): return _is_property(x) and hasattr(x, 'fget') and hasattr(x.fget, 'func')
def _property_getter(x): return x.fget.func if _has_property_getter(x) else x

def _unwrapped_func(x):
    x = _unwrapped_type_dispatch_func(x)
    x = _property_getter(x)
    return x


def get_source_link(func):
    "Return link to `func` in source code"
    import inspect
    func = _unwrapped_func(func)
    try: line = inspect.getsourcelines(func)[1]
    except Exception: return ''
    mod = inspect.getmodule(func)
    module = mod.__name__.replace('.', '/') + '.py'
    try:
        nbdev_mod = import_module(mod.__package__.split('.')[0] + '._modidx')
        ghurl = nested_idx(nbdev_mod.d, 'settings', 'git_url')
        return f"{ghurl}{module}#L{line}"
    except: return f"{module}#L{line}"

# %% ../nbs/03_xtras.ipynb #31a380eb
spark_chars = '▁▂▃▅▆▇'

# %% ../nbs/03_xtras.ipynb #a6371687
def _ceil(x, lim=None): return x if (not lim or x <= lim) else lim

def _sparkchar(x, mn, mx, incr, empty_zero):
    if x is None or (empty_zero and not x): return ' '
    if incr == 0: return spark_chars[0]
    res = int((_ceil(x,mx)-mn)/incr-0.5)
    return spark_chars[res]

# %% ../nbs/03_xtras.ipynb #189818ca
def sparkline(data, mn=None, mx=None, empty_zero=False):
    "Sparkline for `data`, with `None`s (and zero, if `empty_zero`) shown as empty column"
    valid = [o for o in data if o is not None]
    if not valid: return ' '
    mn,mx,n = ifnone(mn,min(valid)),ifnone(mx,max(valid)),len(spark_chars)
    res = [_sparkchar(x=o, mn=mn, mx=mx, incr=(mx-mn)/n, empty_zero=empty_zero) for o in data]
    return ''.join(res)

# %% ../nbs/03_xtras.ipynb #81d193dd
def modify_exception(
    e:Exception, # An exception
    msg:str=None, # A custom message
    replace:bool=False, # Whether to replace e.args with [msg]
) -> Exception:
    "Modifies `e` with a custom message attached"
    e.args = [f'{e.args[0]} {msg}'] if not replace and len(e.args) > 0 else [msg]
    return e

# %% ../nbs/03_xtras.ipynb #6eed0273
def round_multiple(x, mult, round_down=False):
    "Round `x` to nearest multiple of `mult`"
    def _f(x_): return (int if round_down else round)(x_/mult)*mult
    res = L(x).map(_f)
    return res if is_listy(x) else res[0]

# %% ../nbs/03_xtras.ipynb #aa966f19
def set_num_threads(nt):
    "Get numpy (and others) to use `nt` threads"
    try: import mkl; mkl.set_num_threads(nt)
    except: pass
    try: import torch; torch.set_num_threads(nt)
    except: pass
    os.environ['IPC_ENABLE']='1'
    for o in ['OPENBLAS_NUM_THREADS','NUMEXPR_NUM_THREADS','OMP_NUM_THREADS','MKL_NUM_THREADS']:
        os.environ[o] = str(nt)

# %% ../nbs/03_xtras.ipynb #4dac8b41
def join_path_file(file, path, ext=''):
    "Return `path/file` if file is a string or a `Path`, file otherwise"
    if not isinstance(file, (str, Path)): return file
    path.mkdir(parents=True, exist_ok=True)
    return path/f'{file}{ext}'

# %% ../nbs/03_xtras.ipynb #205987ad
def autostart(g):
    "Decorator that automatically starts a generator"
    @functools.wraps(g)
    def f():
        r = g()
        next(r)
        return r
    return f

# %% ../nbs/03_xtras.ipynb #713d7560
class EventTimer:
    "An event timer with history of `store` items of time `span`"

    def __init__(self, store=5, span=60):
        import collections
        self.hist,self.span,self.last = collections.deque(maxlen=store),span,perf_counter()
        self._reset()

    def _reset(self): self.start,self.events = self.last,0

    def add(self, n=1):
        "Record `n` events"
        if self.duration>self.span:
            self.hist.append(self.freq)
            self._reset()
        self.events +=n
        self.last = perf_counter()

    @property
    def duration(self): return perf_counter()-self.start
    @property
    def freq(self): return self.events/self.duration

# %% ../nbs/03_xtras.ipynb #a79aa5d5
_fmt = string.Formatter()

# %% ../nbs/03_xtras.ipynb #b5bc6dd4
def stringfmt_names(s:str)->list:
    "Unique brace-delimited names in `s`"
    return uniqueify(o[1] for o in _fmt.parse(s) if o[1])

# %% ../nbs/03_xtras.ipynb #14ea0571
class PartialFormatter(string.Formatter):
    "A `string.Formatter` that doesn't error on missing fields, and tracks missing fields and unused args"
    def __init__(self):
        self.missing = set()
        super().__init__()

    def get_field(self, nm, args, kwargs):
        try: return super().get_field(nm, args, kwargs)
        except KeyError:
            self.missing.add(nm)
            return '{'+nm+'}',nm

    def check_unused_args(self, used, args, kwargs):
        self.xtra = filter_keys(kwargs, lambda o: o not in used)

# %% ../nbs/03_xtras.ipynb #eb9cb738
def partial_format(s:str, **kwargs):
    "string format `s`, ignoring missing field errors, returning missing and extra fields"
    fmt = PartialFormatter()
    res = fmt.format(s, **kwargs)
    return res,list(fmt.missing),fmt.xtra

# %% ../nbs/03_xtras.ipynb #becf7a2d
def truncstr(s:str, maxlen:int, suf:str='…', space='', sizevar:str=None)->str:
    "Truncate `s` to length `maxlen`, adding suffix `suf` if truncated"
    if sizevar: suf = suf.format_map({sizevar: len(s)})
    return s[:maxlen-len(suf)]+suf if len(s)+len(space)>maxlen else s+space

# %% ../nbs/03_xtras.ipynb #32fa9ba8
def utc2local(dt:datetime)->datetime:
    "Convert `dt` from UTC to local time"
    return dt.replace(tzinfo=timezone.utc).astimezone(tz=None)

# %% ../nbs/03_xtras.ipynb #55abbab7
def local2utc(dt:datetime)->datetime:
    "Convert `dt` from local to UTC time"
    return dt.replace(tzinfo=None).astimezone(tz=timezone.utc)

# %% ../nbs/03_xtras.ipynb #f32ab7a7
def trace(f):
    "Add `set_trace` to an existing function `f`"
    from pdb import set_trace
    if getattr(f, '_traced', False): return f
    def _inner(*args,**kwargs):
        set_trace()
        return f(*args,**kwargs)
    _inner._traced = True
    return _inner

# %% ../nbs/03_xtras.ipynb #45f779a3
@contextmanager
def modified_env(*delete, **replace):
    "Context manager temporarily modifying `os.environ` by deleting `delete` and replacing `replace`"
    prev = dict(os.environ)
    try:
        os.environ.update(replace)
        for k in delete: os.environ.pop(k, None)
        yield
    finally:
        os.environ.clear()
        os.environ.update(prev)

# %% ../nbs/03_xtras.ipynb #12f93aed
class ContextManagers(GetAttr):
    "Wrapper for `contextlib.ExitStack` which enters a collection of context managers"
    def __init__(self, mgrs): self.default,self.stack = L(mgrs),ExitStack()
    def __enter__(self): self.default.map(self.stack.enter_context)
    def __exit__(self, *args, **kwargs): self.stack.__exit__(*args, **kwargs)

# %% ../nbs/03_xtras.ipynb #f58a49b6
def shufflish(x, pct=0.04):
    "Randomly relocate items of `x` up to `pct` of `len(x)` from their starting location"
    n = len(x)
    import random
    return L(x[i] for i in sorted(range_of(x), key=lambda o: o+n*(1+random.random()*pct)))

# %% ../nbs/03_xtras.ipynb #3a527c27
def console_help(libname: str):
    "Show help for all console scripts from `libname`"
    from fastcore.style import S
    from importlib.metadata import entry_points
    for e in entry_points(group='console_scripts'):
        if e.module == libname or e.module.startswith(libname + '.'):
            nm = S.bold.light_blue(e.name)
            print(f'{nm:45}{e.load().__doc__}')

# %% ../nbs/03_xtras.ipynb #01a02b78
def hl_md(s, lang='html', show=True):
    "Syntax highlight `s` using `lang`."
    if hasattr(s, '__html__'): s = s.__html__()
    md = f'```{lang}\n{s}\n```'
    if not show: return md
    try:
        from IPython import display
        return display.Markdown(md)
    except ImportError: print(s)

# %% ../nbs/03_xtras.ipynb #b59609d4
def type2str(typ:type)->str:
    "Stringify `typ`"
    if typ is None or typ is NoneType: return 'None'
    if hasattr(typ, '__origin__'):
        args = ", ".join(type2str(arg) for arg in typ.__args__)
        if typ.__origin__ is Union: return f"Union[{args}]"
        return f"{typ.__origin__.__name__}[{args}]"
    elif isinstance(typ, type): return typ.__name__
    return str(typ)

# %% ../nbs/03_xtras.ipynb #9ef13805
def dataclass_src(cls):
    src = f"@dataclass\nclass {cls.__name__}:\n"
    for f in dataclasses.fields(cls):
        d = "" if f.default is dataclasses.MISSING else f" = {f.default!r}"
        src += f"    {f.name}: {type2str(f.type)}{d}\n"
    return src

# %% ../nbs/03_xtras.ipynb #ecd7f0a2
class Unset(Enum):
    _Unset=''
    def __repr__(self): return 'UNSET'
    def __str__ (self): return 'UNSET'
    def __bool__(self): return False
UNSET = Unset._Unset

# %% ../nbs/03_xtras.ipynb #1e9a46ab
def nullable_dc(cls):
    "Like `dataclass`, but default of `UNSET` added to fields without defaults"
    for k,v in get_annotations_ex(cls)[0].items():
        if not hasattr(cls,k): setattr(cls, k, field(default=UNSET))
    return dataclass(cls)

# %% ../nbs/03_xtras.ipynb #e62e3c77
def make_nullable(clas):
    if hasattr(clas, '_nullable'): return
    clas._nullable = True

    original_init = clas.__init__
    def __init__(self, *args, **kwargs):
        flds = fields(clas)
        dargs = {k.name:v for k,v in zip(flds, args)}
        for f in flds:
            nm = f.name
            if nm not in dargs and nm not in kwargs and f.default is None and f.default_factory is MISSING:
                kwargs[nm] = UNSET
        original_init(self, *args, **kwargs)
    
    clas.__init__ = __init__

    for f in fields(clas):
        if f.default is MISSING and f.default_factory is MISSING: f.default = None
    
    return clas

# %% ../nbs/03_xtras.ipynb #9fa70a25
def flexiclass(
        cls # The class to convert
    ) -> dataclass:
    "Convert `cls` into a `dataclass` like `make_nullable`. Converts in place and also returns the result."
    if is_dataclass(cls): return make_nullable(cls)
    for k,v in get_annotations_ex(cls)[0].items():
        if not hasattr(cls,k) or getattr(cls,k) is MISSING:
            setattr(cls, k, field(default=UNSET))
    return dataclass(cls, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False)

# %% ../nbs/03_xtras.ipynb #7729d860
def asdict(o)->dict:
    "Convert `o` to a `dict`, supporting dataclasses, namedtuples, iterables, and `__dict__` attrs."
    if isinstance(o, dict): return o
    if hasattr(o, '_asdict'): r = o._asdict()
    elif is_dataclass(o): r = dataclasses.asdict(o)
    elif hasattr(o, '__iter__'):
        try: r = dict(o)
        except TypeError: pass
    elif hasattr(o, '__flds__'): r = {x:getattr(o,x) for x in o.__flds__}
    else: r = vars(o)
    skip = set(getattr(o, '__skip__', []))
    return {k:v for k,v in r.items() if v is not UNSET and v is not MISSING and k not in skip}

# %% ../nbs/03_xtras.ipynb #674a899f
def vars_pub(x):
    "Get public non-skipped vars"
    skip = set(getattr(x, '__skip__', []))
    return [o for o in vars(x) if o[0]!='_' and o not in skip]

# %% ../nbs/03_xtras.ipynb #575add44
def is_typeddict(cls:type)->bool:
    "Check if `cls` is a `TypedDict`"
    attrs = 'annotations', 'required_keys', 'optional_keys'
    return isinstance(cls, type) and all(hasattr(cls, f'__{attr}__') for attr in attrs)

# %% ../nbs/03_xtras.ipynb #557c2708
def is_namedtuple(cls):
    "`True` if `cls` is a namedtuple type"
    return issubclass(cls, tuple) and hasattr(cls, '_fields')

# %% ../nbs/03_xtras.ipynb #bb02cce3
class CachedIter:
    "Cache the result returned by an iterator"
    def __init__(self, o): self.o,self.value = o,UNSET
    def __iter__(self):
        if self.value is UNSET: self.value = yield from self.o
        return self.value

# %% ../nbs/03_xtras.ipynb #64471925
class CachedAwaitable:
    "Cache the result from an awaitable"
    def __init__(self, o): self.o,self.value = o,UNSET
    def __await__(self):
        if self.value is UNSET: self.value = yield from self.o.__await__()
        return self.value

# %% ../nbs/03_xtras.ipynb #86cb7a86
def reawaitable(func:callable):
    "Wraps the result of an asynchronous function into an object which can be awaited more than once"
    @wraps(func)
    def _f(*args, **kwargs): return CachedAwaitable(func(*args, **kwargs))
    return _f

# %% ../nbs/03_xtras.ipynb #1fc78492
def is_async_callable(obj):
    "Check if `obj` is an async callable, handling `partial` wrappers and callable instances"
    # Implementation from Starlette; Copyright © 2018, Encode OSS Ltd.
    from asyncio import iscoroutinefunction
    while isinstance(obj, partial): obj = obj.func
    return iscoroutinefunction(obj) or (callable(obj) and iscoroutinefunction(obj.__call__))

# %% ../nbs/03_xtras.ipynb #f7a2e9ed
async def maybe_await(o):
    "Await `o` if needed, and return it"
    from inspect import isawaitable
    return await o if isawaitable(o) else o

# %% ../nbs/03_xtras.ipynb #02f9f070
async def noopa(x=None, *args, **kwargs):
    "Do nothing (async)"
    return x

# %% ../nbs/03_xtras.ipynb #d2b4fe09
def flexicache(*funcs, maxsize=128):
    "Like `lru_cache`, but customisable with policy `funcs`"
    import asyncio
    from collections import OrderedDict
    def _f(func):
        cache,states = OrderedDict(), [None]*len(funcs)
        def _cache_logic(key, execute_func):
            if key in cache:
                result,states = cache[key]
                if not any(f(state) for f,state in zip(funcs, states)):
                    cache[key] = cache.pop(key)
                    return result
                del cache[key]
            try: newres = execute_func()
            except:
                if key not in cache: raise
                cache[key] = cache.pop(key)
                return result
            cache[key] = (newres, [f(None) for f in funcs])
            if len(cache) > maxsize: cache.popitem(last=False)
            return newres

        @wraps(func)
        def wrapper(*args, **kwargs):
            return _cache_logic(f"{args} // {kwargs}", lambda: func(*args, **kwargs))

        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            return await _cache_logic(f"{args} // {kwargs}", lambda: asyncio.ensure_future(func(*args, **kwargs)))

        return async_wrapper if asyncio.iscoroutinefunction(func) else wrapper
    return _f

# %% ../nbs/03_xtras.ipynb #cf131274
def time_policy(seconds):
    "A `flexicache` policy that expires cached items after `seconds` have passed"
    def policy(last_time):
        now = time()
        return now if last_time is None or now-last_time>seconds else None
    return policy

# %% ../nbs/03_xtras.ipynb #8387a8bf
def mtime_policy(filepath):
    "A `flexicache` policy that expires cached items after `filepath` modified-time changes"
    def policy(mtime):
        current_mtime = getmtime(filepath)
        return current_mtime if mtime is None or current_mtime>mtime else None
    return policy

# %% ../nbs/03_xtras.ipynb #1094d811
def timed_cache(seconds=60, maxsize=128):
    "Like `lru_cache`, but also with time-based eviction"
    return flexicache(time_policy(seconds), maxsize=maxsize)
