diff --git a/brainrender/_colors.py b/brainrender/_colors.py index 27b29a8..765b378 100644 --- a/brainrender/_colors.py +++ b/brainrender/_colors.py @@ -1,21 +1,43 @@ +"""Color mapping and palette utilities for brainrender.""" + import random import matplotlib as mpl import numpy as np +import numpy.typing as npt from vedo.colors import colors as vcolors from vedo.colors import get_color as getColor -def map_color(value, name="jet", vmin=None, vmax=None): - """Map a real value in range [vmin, vmax] to a (r,g,b) color scale. +def map_color( + value: float, + name: str = "jet", + vmin: float | None = None, + vmax: float | None = None, +) -> tuple[float, float, float]: + """ + Map a scalar value in ``[vmin, vmax]`` to an RGB colour. + + Parameters + ---------- + value + Scalar value to transform into a colour. + name + Colormap name or matplotlib colormap. Default ``"jet"``. + vmin + Lower bound of the value range. + vmax + Upper bound of the value range. - :param value: scalar value to transform into a color - :type value: float, list - :param name: color map name (Default value = "jet") - :type name: str, matplotlib.colors.LinearSegmentedColorMap - :param vmin: (Default value = None) - :param vmax: (Default value = None) - :returns: return: (r,g,b) color, or a list of (r,g,b) colors. + Returns + ------- + tuple of float + ``(r, g, b)`` colour. + + Raises + ------ + ValueError + If ``vmax`` is smaller than ``vmin``. """ if vmax < vmin: raise ValueError("vmax should be larger than vmin") @@ -31,13 +53,29 @@ def map_color(value, name="jet", vmin=None, vmax=None): return mp(value)[0:3] -def make_palette(N, *colors): - """Generate N colors starting from `color1` to `color2` - by linear interpolation HSV in or RGB spaces. - Adapted from vedo make_palette function +def make_palette(N: int, *colors: str) -> list[npt.NDArray]: + """ + Generate N colours interpolated across the given input colours. + + Adapted from vedo's ``make_palette`` function. Interpolation is + performed in RGB space. + + Parameters + ---------- + N + Number of output colours. + *colors + Input colours. Any number between 1 and N is accepted. + + Returns + ------- + list of numpy.ndarray + List of ``(r, g, b)`` colour arrays. - :param int: N: number of output colors. - :param colors: input colors, any number of colors with 0 < ncolors <= N is okay. + Raises + ------ + ValueError + If no colours are passed or more colours than N are passed. """ N = int(N) @@ -72,9 +110,19 @@ def make_palette(N, *colors): return output -def get_random_colors(n_colors=1): +def get_random_colors(n_colors: int = 1) -> str | list[str]: """ - :param n_colors: (Default value = 1) + Return one or more random colour names from vedo's colour palette. + + Parameters + ---------- + n_colors + Number of colours to return. Default 1. + + Returns + ------- + str or list of str + A single colour name if ``n_colors == 1``, otherwise a list. """ col_names = list(vcolors.keys()) if n_colors == 1: diff --git a/brainrender/_io.py b/brainrender/_io.py index 93d1115..46dfa36 100644 --- a/brainrender/_io.py +++ b/brainrender/_io.py @@ -1,15 +1,34 @@ +"""File I/O and network utilities for brainrender.""" + +from collections.abc import Callable from pathlib import Path +from typing import Any, ParamSpec, TypeVar import requests -from vedo import Mesh, load +from vedo import Mesh, Volume, load + +P = ParamSpec("P") +R = TypeVar("R") -def connected_to_internet(url="http://www.google.com/", timeout=5): +def connected_to_internet( + url: str = "http://www.google.com/", + timeout: int = 5, +) -> bool: """ - Check that there is an internet connection + Check that there is an internet connection. + + Parameters + ---------- + url + URL to use for testing. Default ``"http://www.google.com/"``. + timeout + Timeout in seconds. Default 5. - :param url: url to use for testing (Default value = 'http://www.google.com/') - :param timeout: timeout to wait for [in seconds] (Default value = 5) + Returns + ------- + bool + ``True`` if an internet connection is available, otherwise ``False``. """ try: @@ -20,27 +39,54 @@ def connected_to_internet(url="http://www.google.com/", timeout=5): return False -def fail_on_no_connection(func): +def fail_on_no_connection(func: Callable[P, R]) -> Callable[P, R]: """ - Decorator that throws an error if no internet connection is available + Decorator that raises an error if no internet connection is available. + + Parameters + ---------- + func + Function to wrap. + + Returns + ------- + collections.abc.Callable + + Raises + ------ + ConnectionError + If no internet connection is found. """ if not connected_to_internet(): # pragma: no cover raise ConnectionError( "No internet connection found." ) # pragma: no cover - def inner(*args, **kwargs): + def inner(*args: Any, **kwargs: Any) -> Any: return func(*args, **kwargs) return inner -def request(url): +def request(url: str) -> requests.Response: """ - Sends a request to a url + Send a GET request to a URL. - :param url: + Parameters + ---------- + url + URL to request. + Returns + ------- + requests.Response + + Raises + ------ + ConnectionError + If no internet connection is found. + ValueError + If the request fails. """ if not connected_to_internet(): # pragma: no cover raise ConnectionError( @@ -57,13 +103,29 @@ def request(url): raise ValueError(exception_string) -def check_file_exists(func): # pragma: no cover +def check_file_exists( + func: Callable[P, R], +) -> Callable[P, R]: # pragma: no cover """ - Decorator that throws an error if a function;s first argument + Decorator that raises an error if a function's first argument is not a path to an existing file. + + Parameters + ---------- + func + Function to wrap. + + Returns + ------- + collections.abc.Callable + + Raises + ------ + FileNotFoundError + If the file does not exist. """ - def inner(*args, **kwargs): + def inner(*args: Any, **kwargs: Any) -> Any: if not Path(args[0]).exists(): raise FileNotFoundError( f"File {args[0]} not found" @@ -74,13 +136,26 @@ def inner(*args, **kwargs): @check_file_exists -def load_mesh_from_file(filepath, color=None, alpha=None): +def load_mesh_from_file( + filepath: str | Path, + color: str | None = None, + alpha: float | None = None, +) -> Mesh | Volume: """ - Load a a mesh or volume from files like .obj, .stl, ... + Load a mesh or volume from a file (e.g. .obj, .stl). - :param filepath: path to file - :param **kwargs: + Parameters + ---------- + filepath + Path to the mesh file. + color + Colour to apply to the mesh. + alpha + Transparency to apply to the mesh. + Returns + ------- + vedo.Mesh or vedo.Volume """ actor = load(str(filepath)) actor.c(color).alpha(alpha) diff --git a/brainrender/_utils.py b/brainrender/_utils.py index c20f1c3..947e15d 100644 --- a/brainrender/_utils.py +++ b/brainrender/_utils.py @@ -1,26 +1,55 @@ +"""General utility helpers for file system traversal and list manipulation.""" + from pathlib import Path +from typing import TypeVar + +T = TypeVar("T") -def listdir(fld): +def listdir(fld: str | Path) -> list[str]: """ List the files into a folder with the complete file path instead of the relative file path like os.listdir. - :param fld: string, folder path + Parameters + ---------- + fld + Path to the folder. + Returns + ------- + list of str """ return [str(f) for f in Path(fld).glob("**/*") if f.is_file()] -def get_subdirs(folderpath): +def get_subdirs(folderpath: str | Path) -> list[str]: """ - Returns the subfolders in a given folder + Return all subdirectories in a given folder. + + Parameters + ---------- + folderpath + Path to the folder. + + Returns + ------- + list of str """ return [str(f) for f in Path(folderpath).glob("**/*") if f.is_dir()] -def listify(obj): +def listify(obj: T | list[T] | tuple[T, ...]) -> list[T]: """ - Makes sure that the obj is a list + Ensure the object is a list. + + Parameters + ---------- + obj + Object to listify. + + Returns + ------- + list """ if isinstance(obj, list): return obj @@ -30,11 +59,20 @@ def listify(obj): return [obj] -def return_list_smart(lst): +def return_list_smart(lst: list[T]) -> list[T] | T | None: """ - If the list has length > 1 returns the list - if it has length == 1 it returns the element - if it has length == 0 it returns None + Return a list, single element, or None depending on list length. + + Parameters + ---------- + lst + Input list. + + Returns + ------- + list, Any, or None + The list if length > 1, the single item if length == 1, + or None if empty. """ if len(lst) > 1: return lst diff --git a/brainrender/_video.py b/brainrender/_video.py index 97f8f54..8b0c64a 100644 --- a/brainrender/_video.py +++ b/brainrender/_video.py @@ -1,4 +1,7 @@ +"""Thin wrapper around vedo.Video for rendering brainrender scenes to video.""" + import os +from typing import Any from loguru import logger from myterial import amber_light @@ -7,19 +10,40 @@ class Video(VtkVideo): - # Redefine vedo.Video close method - def __init__(self, *args, fmt="mp4", size="1620x1050", **kwargs): + def __init__( + self, + *args: Any, + fmt: str = "mp4", + size: str = "1620x1050", + **kwargs: Any, + ) -> None: """ - Video class, takes care of storing screenshots (frames) - as images in a temporary folder and then merging these into a - single video file when the video is closed. + Store screenshots as frames and merge them into a video on close. + + Parameters + ---------- + *args + Positional arguments forwarded to ``vedo.Video``. + fmt + Video format. Default ``"mp4"``. + size + Frame size in pixels. Default ``"1620x1050"``. + **kwargs + Keyword arguments forwarded to ``vedo.Video``. """ super().__init__(*args, **kwargs) self.format = fmt self.size = size - def close(self): - """Render the video and write to file.""" + def close(self) -> tuple[int, str]: + """ + Render the video and write to file. + + Returns + ------- + tuple + The FFmpeg exit code and the executed command string. + """ print(f"[{amber_light}]Saving video") logger.debug(f"[{amber_light}]Saving video")