def save()

in lucid/misc/io/saving.py [0:0]

• 52 lines of code
• 21 McCabe index (conditional complexity)

def save(thing, url_or_handle, allow_unsafe_formats=False, save_context: Optional[CaptureSaveContext] = None, **kwargs):
    """Save object to file on CNS.

    File format is inferred from path. Use save_img(), save_npy(), or save_json()
    if you need to force a particular format.

    Args:
      obj: object to save.
      path: CNS path.
      allow_unsafe_formats: set to True to allow saving unsafe formats (eg. pickles)
      save_context: a context into which to capture saves, otherwise will try to use global context

    Raises:
      RuntimeError: If file extension not supported.
    """

    # Determine context
    # Is this a handle? What is the extension? Are we saving to GCS?
    is_handle = hasattr(url_or_handle, "write") and hasattr(url_or_handle, "name")
    if is_handle:
        path = url_or_handle.name
    else:
        path = url_or_handle

    path_without_ext, ext = os.path.splitext(path)
    is_gcs = path.startswith("gs://")

    if ext in compressors:
        compressor = compressors[ext]
        _, ext = os.path.splitext(path_without_ext)
    else:
        compressor = nullcontext

    if not ext:
        raise RuntimeError("No extension in URL: " + path)

    # Determine which saver should be used
    if ext in savers:
        saver = savers[ext]
    elif ext in unsafe_savers:
        if not allow_unsafe_formats:
            raise ValueError(f"{ext} is considered unsafe, you must explicitly allow its use by passing allow_unsafe_formats=True")
        saver = unsafe_savers[ext]
    elif isinstance(thing, str):
        saver = save_str
    else:
        message = "Unknown extension '{}'. As a result, only strings can be saved, not {}. Supported extensions: {}"
        raise ValueError(message.format(ext, type(thing).__name__, list(savers.keys())))

    # Actually save
    if is_handle:
        handle_provider = nullcontext
    else:
        handle_provider = write_handle

    with handle_provider(url_or_handle) as handle:
        with compressor(handle) as compressed_handle:
            result = saver(thing, compressed_handle, **kwargs)

    # Set mime type on gcs if html -- usually, when one saves an html to GCS,
    # they want it to be viewsable as a website.
    if is_gcs and ext == ".html":
        subprocess.run(
            ["gsutil", "setmeta", "-h", "Content-Type: text/html; charset=utf-8", path]
        )
    if is_gcs and ext == ".json":
        subprocess.run(
            ["gsutil", "setmeta", "-h", "Content-Type: application/json", path]
        )

    # capture save if a save context is available
    save_context = save_context if save_context is not None else CaptureSaveContext.current_save_context()
    if save_context:
        log.debug(
            "capturing save: resulted in {} -> {} in save_context {}".format(
                result, path, save_context
            )
        )
        save_context.capture(result)

    if result is not None and "url" in result and result["url"].startswith("gs://"):
        result["serve"] = "https://storage.googleapis.com/{}".format(result["url"][5:])

    return result