Skip to content

Commit

Permalink
fix(io): properly handle registry getattr
Browse files Browse the repository at this point in the history
This fixes an issue where _repr_html_ will fallback to __repr__.
Additionally, `get` will now raise a KeyError instead of a ValueError.
  • Loading branch information
kmnhan committed Mar 25, 2024
1 parent 1fc31af commit 499526f
Showing 1 changed file with 14 additions and 11 deletions.
25 changes: 14 additions & 11 deletions src/erlab/io/dataloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import contextlib
import datetime
import itertools
from multiprocessing import Value
import os
from collections.abc import Iterable, Sequence

Expand Down Expand Up @@ -702,25 +703,27 @@ def register(self, loader_class: type[LoaderBase]):
self.alias_mapping[alias] = loader_class.name

def get(self, key: str) -> LoaderBase:
try:
loader_name = self.alias_mapping.get(key)
loader = self.loaders.get(loader_name)
loader_name = self.alias_mapping.get(key)
loader = self.loaders.get(loader_name)

if not isinstance(loader, LoaderBase):
# If not an instance, create one
self.loaders[loader_name] = loader()
loader = self.loaders[loader_name]
if loader is None:
raise KeyError(f"Loader for {key} not found")

return loader
if not isinstance(loader, LoaderBase):
# If not an instance, create one
self.loaders[loader_name] = loader()
loader = self.loaders[loader_name]

except KeyError:
raise ValueError(f"Loader for {key} not found")
return loader

def __getitem__(self, key: str) -> LoaderBase:
return self.get(key)

def __getattr__(self, key: str) -> LoaderBase:
return self.get(key)
try:
return self.get(key)
except KeyError:
raise AttributeError(f"Loader for {key} not found")

def set_loader(self, loader: str | LoaderBase):
"""Set the current data loader.
Expand Down

0 comments on commit 499526f

Please sign in to comment.