core ¤

def add_related_collection(
    self,
    track_ids: List[Union[str, uuid.UUID]],
    name: str,
    description: str = "",
    user_id: Optional[Union[str, uuid.UUID]] = None,
    relationship_type: str = "relationship",
    meta: Optional[Dict[str, Any]] = None,
) -> NendoCollection:
    """Create a new collection with a relationship to the current collection.

    Args:
        track_ids (List[Union[str, uuid.UUID]]): List of track ids.
        name (str): Name of the new related collection.
        description (str): Description of the new related collection.
        user_id (UUID, optional): The ID of the user adding the collection.
        relationship_type (str): Type of the relationship.
        meta (Dict[str, Any]): Meta of the new related collection.

    Returns:
        schema.NendoCollection: The newly added NendoCollection object.
    """
    self.nendo_instance.library.add_related_collection(
        track_ids=track_ids,
        collection_id=self.id,
        name=name,
        description=description,
        user_id=user_id,
        relationship_type=relationship_type,
        meta=meta,
    )
    self.refresh()
    return self

def add_track(
    self,
    track_id: Union[str, uuid.UUID],
    position: Optional[int] = None,
    meta: Optional[Dict[str, Any]] = None,
) -> NendoCollection:
    """Creates a relationship from the track to the collection.

    Args:
        track_id (Union[str, uuid.UUID]): ID of the track to add.
        position (int, optional): Target position of the track inside
            the collection.
        meta (Dict[str, Any]): Metadata of the relationship.

    Returns:
        NendoCollection: The collection itself.
    """
    updated_collection = self.nendo_instance.library.add_track_to_collection(
        track_id=track_id,
        collection_id=self.id,
        position=position,
        meta=meta,
    )
    self.related_tracks = updated_collection.related_tracks
    return self

def delete(
    self,
    remove_relationships: bool = False,
) -> NendoCollection:
    """Deletes the collection from the nendo library.

    Args:
        remove_relationships (bool, optional):
            If False prevent deletion if related tracks exist,
            if True delete relationships together with the object.
            Defaults to False.

    Returns:
        NendoCollection: The collection itself.
    """
    self.nendo_instance.library.remove_collection(
        collection_id=self.id,
        remove_relationships=remove_relationships,
    )
    return self

def export(
    self,
    export_path: str,
    filename_suffix: str = "nendo",
    file_format: str = "wav",
) -> NendoCollection:
    """Export the collection to a directory.

    Args:
        export_path (str): Path to a directory into which the collection's tracks
            should be exported.
        filename_suffix (str): The suffix which should be appended to each
            exported track's filename.
        file_format (str, optional): Format of the exported track. Ignored if
            file_path is a full file path. Defaults to "wav".

    Returns:
        NendoTrack: The track itself.
    """
    self.nendo_instance.library.export_collection(
        collection_id=self.id,
        export_path=export_path,
        filename_suffix=filename_suffix,
        file_format=file_format,
    )
    return self

def get_meta(self, key: str) -> Dict[str, Any]:
    """Get the meta entry for the given key.

    Args:
        key (str): The key to get metadata for.

    Returns:
        Dict[str, Any]: Meta entry for given key.
    """
    if not self.has_meta(key):
        logger.warning("Key not found in meta: %s", key)
        return None
    return self.meta[key]

def get_related_collections(
    self,
    direction: str = "to",
    user_id: Optional[Union[str, uuid.UUID]] = None,
    order_by: Optional[str] = None,
    order: Optional[str] = "asc",
    limit: Optional[int] = None,
    offset: Optional[int] = None,
) -> List[NendoCollection]:
    """Get all collections to which the current collection has a relationship.

    Args:
        direction (str, optional): The relationship direction.
            Can be either one of "to", "from", or "both".
            Defaults to "to".
        user_id (Union[str, UUID], optional): The user ID to filter for.
        order_by (Optional[str]): Key used for ordering the results.
        order (Optional[str]): Order in which to retrieve results ("asc" or "desc").
        limit (Optional[int]): Limit the number of returned results.
        offset (Optional[int]): Offset into the paginated results (requires limit).

    Returns:
        List[NendoCollection]: List containting all related NendoCollections
    """
    return self.nendo_instance.library.get_related_collections(
        collection_id=self.id,
        direction=direction,
        user_id=user_id,
        order_by=order_by,
        order=order,
        limit=limit,
        offset=offset,
    )

def has_meta(self, key: str) -> bool:
    """Check if a given collection has the given key in its meta dict.

    Args:
        key (str): The key to check for.

    Returns:
        bool: True if the key exists, False otherwise.
    """
    return any(k == key for k in self.meta)

def has_related_collection(
    self,
    collection_id: Union[str, uuid.UUID],
) -> bool:
    """Check if the collection has a relationship to the specified collection ID.

    Args:
        collection_id (Union[str, uuid.UUID]): The target collection ID for which
            to check if it is related to the current collection.

    Returns:
        bool: True if the collection has a relationship either from or to the
            collection with the given `collection_id`. False otherwise.
    """
    collection_id = ensure_uuid(collection_id)
    related_collections = self.get_related_collections(
        direction="both",
    )
    return any(rt.id == collection_id for rt in related_collections)

def has_relationship(self, relationship_type: str = "relationship") -> bool:
    """Check if the collection has the specified relationship type.

    Args:
        relationship_type (str): Type of the relationship to check for.
            Defaults to "relationship".

    Returns:
        bool: True if a relationship of the given type exists, False otherwise.
    """
    if self.related_collections is None:
        return False

    return any(
        r.relationship_type == relationship_type for r in self.related_collections
    )

def process(self, plugin: str, **kwargs: Any) -> NendoCollection:
    """Process the collection with the specified plugin.

    Args:
        plugin (str): Name of the plugin to run on the collection.

    Returns:
        NendoCollection: The collection that was created by the plugin.
    """
    registered_plugin: RegisteredNendoPlugin = getattr(
        self.nendo_instance.plugins,
        plugin,
    )
    wrapped_method = get_wrapped_methods(registered_plugin.plugin_instance)

    if len(wrapped_method) > 1:
        raise NendoError(
            "Plugin has more than one wrapped method. Please use `nd.plugins.<plugin_name>.<method_name>` instead.",
        )
    return getattr(self.nendo_instance.plugins, plugin)(collection=self, **kwargs)

def refresh(self) -> NendoCollection:
    """Refreshes the collection with the latest values from the library.

    Returns:
        NendoCollection: The collection itself.
    """
    refreshed_collection = self.nendo_instance.library.get_collection(self.id)
    for attr in vars(refreshed_collection):
        setattr(self, attr, getattr(refreshed_collection, attr))
    return self

def remove_meta(self, key: str) -> NendoCollection:
    """Remove the meta entry for the given key.

    Args:
        key (str): The key to remove metadata for.

    Returns:
        NendoCollection: The track itself.
    """
    if not self.has_meta(key):
        logger.error("Key not found in meta: %s", key)
        return None
    _ = self.meta.pop(key, None)
    self.nendo_instance.library.update_collection(collection=self)
    return self

def remove_track(
    self,
    track_id: Union[str, uuid.UUID],
) -> NendoCollection:
    """Removes the track from the collection.

    Args:
        track_id (Union[str, uuid.UUID]): ID of the track to remove
            from the collection.

    Returns:
        NendoCollection: The collection itself.
    """
    self.nendo_instance.library.remove_track_from_collection(
        track_id=track_id,
        collection_id=self.id,
    )
    self.refresh()
    return self

def save(self) -> NendoCollection:
    """Save the collection to the nendo library.

    Returns:
        NendoCollection: The collection itself.
    """
    self.nendo_instance.library.update_collection(collection=self)
    return self

def set_meta(self, meta: Dict[str, Any]) -> NendoCollection:
    """Set metadata of collection.

    Args:
        meta (Dict[str, Any]): Dictionary containing the metadata to be set.

    Returns:
        NendoCollection: The collection itself.
    """
    try:
        self.meta.update(meta)
        self.nendo_instance.library.update_collection(collection=self)
    except NendoError as e:
        logger.exception("Error updating meta: %s", e)
    return self

def tracks(self) -> List[NendoTrack]:
    """Return all tracks listed in the collection.

    Collection will be loaded from the DB if not already loaded.

    Returns:
        List[NendoTrack]: List of tracks.
    """
    ts = self.__dict__.get("loaded_tracks")
    if ts is None:
        ts = self.nendo_instance.library.get_collection_tracks(self.id)
        self.__dict__["loaded_tracks"] = ts
    return ts

@staticmethod
def batch_process(func):
    """Decorator to run functions multithreaded in batches.

    This decorator function transforms the given function to run
    in multiple threads. It expects that the first argument to the function
    is a list of items, which will be processed in parallel,
    in batches of a given size.
    """

    @functools.wraps(func)
    def wrapper(self, track=None, file_paths=None, *args, **kwargs):
        target = track or file_paths
        if isinstance(target, NendoTrack):
            return func(self, track=target, **kwargs)
        elif isinstance(target, list):  # noqa: RET505
            max_threads = self.config.max_threads
            batch_size = self.config.batch_size
            total = len(target)
            batches = [
                target[i : i + batch_size] for i in range(0, total, batch_size)
            ]
            start_time = time.time()
            futures = []

            def run_batch(batch_index, batch):
                try:
                    batch_start_time = time.time()
                    results = []
                    if track:
                        for _, item in enumerate(batch):
                            result = func(
                                self,
                                track=item,
                                *args,  # noqa: B026
                                **kwargs,
                            )
                            results.extend(result)
                    elif file_paths:
                        result = func(
                            self,
                            file_paths=batch,
                            *args,  # noqa: B026
                            **kwargs,
                        )
                        results.extend(result)
                    batch_end_time = time.time()
                    batch_time = time.strftime(
                        "%H:%M:%S",
                        time.gmtime(batch_end_time - batch_start_time),
                    )
                    total_elapsed_time = batch_end_time - start_time
                    average_time_per_batch = total_elapsed_time / (batch_index + 1)
                    estimated_total_time = average_time_per_batch * len(batches)
                    estimated_total_time_print = time.strftime(
                        "%H:%M:%S",
                        time.gmtime(estimated_total_time),
                    )
                    remaining_time = time.strftime(
                        "%H:%M:%S",
                        time.gmtime(estimated_total_time - total_elapsed_time),
                    )
                    logger.info(
                        f"Finished batch {batch_index + 1}/{len(batches)}.\n"
                        f"Time taken for this batch: {batch_time} - "
                        f"Estimated total time: {estimated_total_time_print} - "
                        f"Estimated remaining time: {remaining_time}\n",
                    )
                    return results
                except NendoError as e:
                    logger.exception(
                        "Error processing batch %d: %s",
                        batch_index,
                        e,
                    )

            with ThreadPoolExecutor(max_workers=max_threads) as executor:
                for batch_index, batch in enumerate(batches):
                    futures.append(executor.submit(run_batch, batch_index, batch))

            all_results = []
            for future in as_completed(futures):
                result = future.result()
                if result:
                    all_results.extend(future.result())
            return all_results
        else:
            raise TypeError("Expected NendoTrack or list of NendoTracks")

    return wrapper

@staticmethod
def stream_output(func):
    """Decorator to turn on streaming mode for functions.

    The requirement for this decorator to work on a function is that it would
    normally return a list.
    """

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        result = func(self, *args, **kwargs)
        if self.config.stream_mode:
            return result
        # function is yielding single tracks if stream_chunk_size == 1
        elif self.config.stream_chunk_size > 1:  # noqa: RET505
            return [track for chunk in result for track in chunk]
        else:
            return list(result)

    return wrapper

def add(
    self,
    plugin_name: str,
    version: str,
    plugin_instance: NendoPlugin,
) -> RegisteredNendoPlugin:
    """Add a Registered plugin to the plugin registry.

    Args:
        plugin_name (str): Name of the plugin.
        version (str): Version of the plugin.
        plugin_instance (schema.NendoPlugin): Instantiated plugin class.

    Returns:
        RegisteredNendoPlugin: The registered nendo plugin.
    """
    self._plugins[plugin_name] = RegisteredNendoPlugin(
        name=plugin_name,
        version=version,
        plugin_instance=plugin_instance,
    )

def all_names(self) -> List[str]:
    """Return all plugins that are registered as a list of their names.

    Returns:
        List[str]: List containing all plugin names.
    """
    return [k for k, v in self._plugins.items()]

def remove(self, plugin_name: str) -> None:
    """Remove a plugin from the plugin registry.

    Args:
        plugin_name (str): Name of the plugin to remove.
    """
    del self._plugins[plugin_name]

@abstractmethod
def as_local(self, file_path: str, location: ResourceLocation, user_id: str) -> str:
    """Get a local handle to the file.

    Return a file given as `file_path` as a locally accessible file,
    even if it originally resides in a remote storage location

    Args:
        file_path (str): Name of the file to obtain a local copy of.
        location (ResourceLocation): Location of the file, as provided by the
            corresponding NendoResource.
        user_id (str): User ID

    Returns:
        str: Local path to the original file (if it was local) or to a
            temporary copy (if it was remote)
    """
    raise NotImplementedError

@abstractmethod
def file_exists(self, file_name: str, user_id: str) -> bool:
    """Check if the given file_name exists in the storage.

    Args:
        file_name (str): Name of the file for which to check existence
        user_id (str): User ID

    Returns:
        bool: True if it exists, false otherwise
    """
    raise NotImplementedError

@abstractmethod
def generate_filename(self, filetype: str, user_id: str) -> str:
    """Generate a collision-free new filename.

    Args:
        filetype (str): The filetype to append (without the dot).
        user_id (str): The ID of the user requesting the file.

    Returns:
        str: The generated filename.
    """
    raise NotImplementedError

@abstractmethod
def get_bytes(self, file_name: str, user_id: str) -> Any:
    """Load the data bytes from the storage.

    Loading includes unpickling the file given as `file_name`.

    Args:
        file_name (str): Name of the target file to load.
        user_id (str): The ID of the user writing the file.

    Returns:
        The deserialized data bytes
    """
    raise NotImplementedError

@abstractmethod
def get_checksum(self, file_name: str, user_id: str) -> str:
    """Compute the checksum for the given file and user_id.

    Args:
        file_name (str): The name of the file in the library for which
            to compute the checksum.
        user_id (str): The ID of the user requesting the checksum.

    Returns:
        str: The checksum of the target file.
    """
    raise NotImplementedError

@abstractmethod
def get_driver_location(self) -> str:
    """Get the default resource location of the storage driver.

    e.g. "original", "local", "gcs", "s3", etc.

    Returns:
        str: Location type.
    """

@abstractmethod
def get_file(self, file_name: str, user_id: str) -> str:
    """Obtains the full path to the file by the name of file_name from storage.

    Args:
        file_name (str): The name of the target file.
        user_id (str): The ID of the user requesting the file.

    Returns:
        str: The full path to the file.
    """
    raise NotImplementedError

@abstractmethod
def get_file_name(self, src: str, user_id: str) -> str:
    """Returns the filename of a resource.

    Args:
        src (str): The full path to the resource.
        user_id (str): The ID of the user requesting the path.

    Returns:
        str: The resource file name, minus the path.
    """
    raise NotImplementedError

@abstractmethod
def get_file_path(self, src: str, user_id: str) -> str:
    """Returns the path of a resource.

    Args:
        src (str): The full path to the resource.
        user_id (str): The ID of the user requesting the path.

    Returns:
        str: The resource path, minus the file name.
    """
    raise NotImplementedError

@abstractmethod
def init_storage_for_user(self, user_id: str) -> Any:
    """Initializes the storage location for the given user.

    Args:
        user_id (str): ID of the user for whom the storage is to be initialized.

    Returns:
        Any: Storage object
    """
    raise NotImplementedError

@abstractmethod
def list_files(self, user_id: str) -> List[str]:
    """Lists all files found in the user's library.

    Args:
        user_id (str): The ID of the user.

    Returns:
        List[str]: List of paths to files.
    """
    raise NotImplementedError

@abstractmethod
def remove_file(self, file_name: str, user_id: str) -> bool:
    """Remove the file given by file_name and user_id from the storage.

    Args:
        file_name (str): Name of the file to remove.
        user_id (str): ID of the user requesting the removal.

    Returns:
        bool: True if deletion was successful, False otherwise.
    """
    raise NotImplementedError

@abstractmethod
def save_bytes(self, file_name: str, data: bytes, user_id: str) -> str:
    """Save a data given as bytes to the FS by pickling them first.

    Args:
        file_name (str): Name of the target file to save the data to.
        data (np.ndarray): The data to write to file in pickled form.
        user_id (str): The ID of the user writing the file.

    Returns:
        str: The full path to the saved file resource.
    """
    raise NotImplementedError

@abstractmethod
def save_file(self, file_name: str, file_path: str, user_id: str) -> str:
    """Saves the file with the given file_name to the path specified by file_path.

    Args:
        file_name (str): The name of the target file.
        file_path (str): The path to the target file.
        user_id (str): The ID of the user requesting the file.

    Returns:
        str: The full path to the saved file resource.
    """
    raise NotImplementedError

@abstractmethod
def save_signal(
    self,
    file_name: str,
    signal: np.ndarray,
    sr: int,
    user_id: str,
) -> str:
    """Save a signal given as a numpy array to storage.

    Args:
        file_name (str): Name of the target file to save the signal to.
        signal (np.ndarray): The signal to write as a numpy array.
        sr (int): The sample rate of the signal.
        user_id (str): The ID of the user writing the file.

    Returns:
        str: The full path to the saved file resource.
    """
    raise NotImplementedError

def as_local(self, file_path: str, location: ResourceLocation, user_id: str) -> str:
    """Get a local handle to the file."""
    if location == self.get_driver_location():
        return self.get_file(os.path.basename(file_path), user_id)
    return file_path

def file_exists(self, file_name: str, user_id: str) -> bool:
    """Check if the given file exists."""
    return os.path.isfile(self.get_file(file_name, user_id))

def generate_filename(self, filetype: str, user_id: str) -> str:  # noqa: ARG002
    """Generate a unique filename."""
    return f"{uuid.uuid4()!s}.{filetype}"

def get_bytes(self, file_name: str, user_id: str) -> Any:
    """Get bytes from a stored file by unpickling it."""
    target_file_path = self.get_file(file_name=file_name, user_id=user_id)
    with open(target_file_path, "rb") as target_file:
        return pickle.loads(target_file.read())  # noqa: S301

def get_checksum(self, file_name: str, user_id: str) -> str:
    """Compute the MD5 checksum of the given file."""
    return md5sum(self.get_file(file_name=file_name, user_id=user_id))

def get_driver_location(self) -> ResourceLocation:
    """Get the default resource location of the storage driver."""
    return ResourceLocation.local

def get_file(self, file_name: str, user_id: str) -> str:
    """Get the full path to the file."""
    return os.path.join(self.library_path, user_id, file_name)

def get_file_name(self, src: str, user_id: str) -> str:  # noqa: ARG002
    """Get the file name (without the path)."""
    return os.path.basename(src)

def get_file_path(self, src: str, user_id: str) -> str:  # noqa: ARG002
    """Get the path to the file (without the file name)."""
    return os.path.dirname(src)

def init_storage_for_user(self, user_id: str) -> str:
    """Initialize local storage for user."""
    user_lib_path = os.path.join(self.library_path, user_id)
    if not os.path.isdir(user_lib_path):
        logger.info(
            f"Library path {user_lib_path} does not exist, creating now.",
        )
        os.makedirs(user_lib_path)
    return user_lib_path

def list_files(self, user_id: str) -> List[str]:
    """List all files contained in the storage."""
    with os.scandir(os.path.join(self.library_path, user_id)) as entries:
        return [entry.name for entry in entries if entry.is_file()]

def remove_file(self, file_name: str, user_id: str) -> bool:
    """Remove the given file."""
    target_file = self.get_file(file_name=file_name, user_id=user_id)
    try:
        os.remove(target_file)
        return True
    except OSError as e:
        logger.error("Removing file %s failed: %s", target_file, e)
        return False

def save_bytes(
    self,
    file_name: str,
    data: bytes,
    user_id: str,
) -> str:
    """Save the given bytes to storage."""
    target_file_path = self.get_file(file_name=file_name, user_id=user_id)
    with open(target_file_path, "wb") as target_file:
        pickle.dump(data, target_file)
    return target_file_path

def save_file(self, file_name: str, file_path: str, user_id: str) -> str:
    """Copy the source file given by file_path to the library."""
    target_file = self.get_file(file_name=file_name, user_id=user_id)
    shutil.copy2(file_path, target_file)
    return target_file

def save_signal(
    self,
    file_name: str,
    signal: np.ndarray,
    sr: int,
    user_id: str,
) -> str:
    """Save the given signal to storage."""
    target_file_path = self.get_file(file_name=file_name, user_id=user_id)
    sf.write(target_file_path, signal, sr, subtype="PCM_16")
    return target_file_path

def add_plugin_data(
    self,
    key: str,
    value: str,
    plugin_name: str,
    plugin_version: Optional[str] = None,
    user_id: Optional[Union[str, uuid.UUID]] = None,
    replace: Optional[bool] = None,
) -> NendoTrack:
    """Add plugin data to a NendoTrack and persist changes into the DB.

    Args:
        key (str): Key under which to save the data.
        value (Any): Data to save.
        plugin_name (str): Name of the plugin.
        plugin_version (str): Version of the plugin. If none is given,
            the currently loaded version of the plugin given by `plugin_name`
            will be used.
        user_id (Union[str, UUID], optional): ID of user adding the plugin data.
        replace (bool, optional): Flag that determines whether
            the last existing data point for the given plugin name and -version
            is overwritten or not. If undefined, the nendo configuration's
            `replace_plugin_data` value will be used.

    Returns:
        NendoTrack: The track itself.
    """
    pd = self.nendo_instance.library.add_plugin_data(
        track_id=self.id,
        plugin_name=plugin_name,
        plugin_version=plugin_version,
        key=key,
        value=value,
        user_id=user_id,
        replace=replace,
    )
    self.plugin_data.append(pd)
    return self