Coverage for src/importnb/loaders.py: 90%
39 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-04 11:11 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-04 11:11 -0800
1from .loader import Loader, SourceModule
2from dataclasses import dataclass, field
3from types import ModuleType
5class DataModule(SourceModule):
6 def _repr_json_(self):
7 return self.data, dict(root=repr(self), expanded=False)
10@dataclass
11class DataStreamLoader(Loader):
12 """an import loader for data streams"""
13 module_type: ModuleType = field(default=DataModule)
15 def exec_module(self, module):
16 with open(module.__file__, "rb") as file:
17 module.data = self.get_data_loader()(file)
18 return module
20 def get_data_loader(self):
21 raise NotImplementedError("load_data not implemented.")
24@dataclass
25class Json(DataStreamLoader):
26 """an import loader for json files"""
28 extensions: tuple = field(default_factory=[".json"].copy)
30 def get_data_loader(self):
31 from json import load
33 return load
36@dataclass
37class Yaml(DataStreamLoader):
38 """an import loader for yml and yaml"""
40 extensions: tuple = field(default_factory=[".yml", ".yaml"].copy)
42 def get_data_loader(self):
43 try:
44 from ruamel.yaml import safe_load
45 except ModuleNotFoundError:
46 from yaml import safe_load
47 # probably want an error message about how to fix this if we cant find yamls
48 return safe_load
51@dataclass
52class Toml(DataStreamLoader):
53 """an import loader for toml"""
55 extensions: tuple = field(default_factory=[".toml"].copy)
57 def get_data_loader(self):
58 try:
59 from tomllib import load
60 except ModuleNotFoundError:
61 from tomli import load
62 return load