# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01v_utils.wandb.ipynb.

# %% auto 0
__all__ = ['WandbRun', 'wandb_upload']

# %% ../../nbs/01v_utils.wandb.ipynb 2
from fastcore.all import *
import polvo as pv

# %% ../../nbs/01v_utils.wandb.ipynb 3
class WandbRun:
    def __init__(self, run=None, **kwargs):
        "Context manager that finishes run when exiting"
        import wandb
        self.run = run or wandb.init(**kwargs)
    def __enter__(self): return self.run
    def __exit__(self, exc_type, exc_value, traceback): self.run.finish()

# %% ../../nbs/01v_utils.wandb.ipynb 4
def wandb_upload(artifact_name:str, *path:str, project=None, type='dataset', run=None):
    "Uploads files or dir to wandb"
    import wandb
    run = run or wandb.init(project=project, job_type="data-upload")
    artifact = wandb.Artifact(artifact_name, type=type)
    
    def _add_file(p): artifact.add_file(p, p)
    for p in path:
        if Path(p).is_file(): _add_file(p)
        else:
            for f in pv.get_files(p): _add_file(f)

    run.log_artifact(artifact)
