diff --git a/src/guidellm/extras/audio.pyi b/src/guidellm/extras/audio.pyi index d1e7e1ee7..652bc3bc5 100644 --- a/src/guidellm/extras/audio.pyi +++ b/src/guidellm/extras/audio.pyi @@ -1,3 +1,5 @@ -from torchcodec import AudioSamples as AudioSamples -from torchcodec.decoders import AudioDecoder as AudioDecoder -from torchcodec.encoders import AudioEncoder as AudioEncoder +from torchcodec import AudioSamples +from torchcodec.decoders import AudioDecoder +from torchcodec.encoders import AudioEncoder + +__all__ = ["AudioDecoder", "AudioEncoder", "AudioSamples"] diff --git a/src/guidellm/extras/vision.py b/src/guidellm/extras/vision.py index 06240e34d..1ce721aed 100644 --- a/src/guidellm/extras/vision.py +++ b/src/guidellm/extras/vision.py @@ -5,9 +5,9 @@ __getattr__, __dir__, __all__ = lazy.attach_extras( __name__, attrs={ - "PILImage": lazy.ExtraAttr("PIL", alias="Image"), - "Image": lazy.ExtraAttr("PIL.Image", alias="Image"), - "iio": lazy.ExtraAttr("imageio", alias="v3"), + "PILImage": lazy.ExtraAttr("PIL", "Image"), + "Image": lazy.ExtraAttr("PIL.Image", "Image"), + "iio": lazy.ExtraAttr("imageio", "v3"), }, error_message="Please install guidellm[vision] to use image/video features", ) diff --git a/src/guidellm/extras/vision.pyi b/src/guidellm/extras/vision.pyi index ce2393bfe..64f4e7fab 100644 --- a/src/guidellm/extras/vision.pyi +++ b/src/guidellm/extras/vision.pyi @@ -1,7 +1,5 @@ import imageio.v3 as iio -from PIL import Image as _PILImage -from PIL.Image import Image as Image +from PIL import Image as PILImage +from PIL.Image import Image __all__ = ["Image", "PILImage", "iio"] - -PILImage = _PILImage diff --git a/src/guidellm/extras/vllm.pyi b/src/guidellm/extras/vllm.pyi index 175bbee52..8deb05ca9 100644 --- a/src/guidellm/extras/vllm.pyi +++ b/src/guidellm/extras/vllm.pyi @@ -1,4 +1,8 @@ -from vllm import AsyncEngineArgs as AsyncEngineArgs -from vllm import AsyncLLMEngine as AsyncLLMEngine -from vllm import RequestOutput as RequestOutput -from vllm import SamplingParams as SamplingParams +from vllm import AsyncEngineArgs, AsyncLLMEngine, RequestOutput, SamplingParams + +__all__ = [ + "AsyncEngineArgs", + "AsyncLLMEngine", + "RequestOutput", + "SamplingParams", +] diff --git a/src/guidellm/utils/lazy_loader.py b/src/guidellm/utils/lazy_loader.py index e5503e889..f0a1e9bc3 100644 --- a/src/guidellm/utils/lazy_loader.py +++ b/src/guidellm/utils/lazy_loader.py @@ -56,13 +56,13 @@ class ExtraAttr(typing.NamedTuple): """Descriptor for a lazily imported attribute in :func:`attach_extras`. - :param source: Dotted module path to import from. - :param alias: Attribute name inside *source*. When ``None`` (the - default), the dictionary key passed to ``attach_extras`` is used. + :param path: Dotted module path to import from. + :param name: Attribute name inside *path*. Only nessary if the attribute name + differs from the exported name. """ - source: str - alias: str | None = None + path: str + name: str | None = None threadlock = threading.Lock() @@ -160,11 +160,11 @@ def __dir__(): return __getattr__, __dir__, __all__.copy() -def _attach_extras_attrs(module_name, attrs, error_message): +def _attach_extras_attrs(module_name, attrs: dict[str, ExtraAttr], error_message): _attr_map = {} for export_name, spec in attrs.items(): - source_attr = spec.alias if spec.alias is not None else export_name - _attr_map[export_name] = (spec.source, source_attr) + source_attr = spec.name if spec.name is not None else export_name + _attr_map[export_name] = (spec.path, source_attr) _all = sorted(_attr_map.keys()) @@ -242,7 +242,7 @@ def attach_extras( :param module_name: Typically use ``__name__``. :param attrs: Map of exported names to :class:`ExtraAttr` descriptors. - Each value specifies the *source* module and an optional *alias* + Each value specifies the *source* module and an optional *attr* (the attribute name inside *source*, when it differs from the dictionary key). :param package: Name of a package whose public attributes should be diff --git a/tests/unit/utils/test_lazy_loader.py b/tests/unit/utils/test_lazy_loader.py index 981c5af97..00262a820 100644 --- a/tests/unit/utils/test_lazy_loader.py +++ b/tests/unit/utils/test_lazy_loader.py @@ -299,7 +299,7 @@ def test_attach_extras_attrs_installed_package(): "test_extras_attrs", attrs={ "sin": lazy.ExtraAttr("math"), - "mypi": lazy.ExtraAttr("math", alias="pi"), + "mypi": lazy.ExtraAttr("math", "pi"), }, error_message="install math", ) @@ -313,7 +313,7 @@ def test_attach_extras_attrs_alias(): """## WRITTEN BY AI ##""" ga, _, _ = lazy.attach_extras( "test_extras_alias", - attrs={"my_sep": lazy.ExtraAttr("os.path", alias="sep")}, + attrs={"my_sep": lazy.ExtraAttr("os.path", "sep")}, error_message="install os", ) import os.path @@ -325,7 +325,7 @@ def test_attach_extras_attrs_submodule_fallback(): """## WRITTEN BY AI ##""" ga, _, _ = lazy.attach_extras( "test_extras_submod", - attrs={"path": lazy.ExtraAttr("os", alias="path")}, + attrs={"path": lazy.ExtraAttr("os", "path")}, error_message="install os", ) import os.path