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

1from .loader import Loader, SourceModule 

2from dataclasses import dataclass, field 

3from types import ModuleType 

4 

5class DataModule(SourceModule): 

6 def _repr_json_(self): 

7 return self.data, dict(root=repr(self), expanded=False) 

8 

9 

10@dataclass 

11class DataStreamLoader(Loader): 

12 """an import loader for data streams""" 

13 module_type: ModuleType = field(default=DataModule) 

14 

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 

19 

20 def get_data_loader(self): 

21 raise NotImplementedError("load_data not implemented.") 

22 

23 

24@dataclass 

25class Json(DataStreamLoader): 

26 """an import loader for json files""" 

27 

28 extensions: tuple = field(default_factory=[".json"].copy) 

29 

30 def get_data_loader(self): 

31 from json import load 

32 

33 return load 

34 

35 

36@dataclass 

37class Yaml(DataStreamLoader): 

38 """an import loader for yml and yaml""" 

39 

40 extensions: tuple = field(default_factory=[".yml", ".yaml"].copy) 

41 

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 

49 

50 

51@dataclass 

52class Toml(DataStreamLoader): 

53 """an import loader for toml""" 

54 

55 extensions: tuple = field(default_factory=[".toml"].copy) 

56 

57 def get_data_loader(self): 

58 try: 

59 from tomllib import load 

60 except ModuleNotFoundError: 

61 from tomli import load 

62 return load 

63