Coverage for src/configuraptor/loaders/register.py: 100%

27 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-21 10:19 +0200

1""" 

2Exposes `register_loader` to define loader for specific file types. 

3""" 

4 

5import typing 

6from pathlib import Path 

7 

8from ._types import T_config 

9 

10T_loader = typing.Callable[[typing.BinaryIO, Path], T_config] 

11T_WrappedLoader = typing.Callable[[T_loader], T_loader] 

12 

13LOADERS: dict[str, T_loader] = {} 

14 

15 

16@typing.overload 

17def register_loader(*extension_args: str) -> T_WrappedLoader: 

18 """ 

19 Overload for case with parens. 

20 

21 @register_loader("yaml", ".yml") 

22 def load_yaml(...): 

23 ... 

24 

25 # extension_args is a tuple of strings 

26 # this will return a wrapper which takes `load_yaml` as input and output. 

27 """ 

28 

29 

30@typing.overload 

31def register_loader(*extension_args: T_loader) -> T_loader: 

32 """ 

33 Overload for case without parens. 

34 

35 @register_loader 

36 def json(...): 

37 ... 

38 

39 # extension_args is a tuple of 1: `def json` 

40 # this will simply return the `json` method itself. 

41 """ 

42 

43 

44def register_loader(*extension_args: str | T_loader) -> T_loader | T_WrappedLoader: 

45 """ 

46 Register a data loader for a new filetype. 

47 

48 Used as a decorator on a method that takes two arguments: 

49 (BinaryIO, Path) - an open binary file stream to the config file and the pathlib.Path to the config file. 

50 By default, the open file handler can be used. 

51 However, some loaders (such as .ini) don't support binary file streams. 

52 These can use the Path to open and read the file themselves however they please. 

53 """ 

54 f_outer = None 

55 extension_set = set() 

56 

57 for extension in extension_args: 

58 if not isinstance(extension, str): 

59 f_outer = extension 

60 extension = extension.__name__ 

61 

62 elif extension.startswith("."): 

63 extension = extension.removeprefix(".") 

64 

65 extension_set.add(extension) 

66 

67 def wrapper(f_inner: T_loader) -> T_loader: 

68 LOADERS.update({ext: f_inner for ext in extension_set}) 

69 return f_inner 

70 

71 if f_outer: 

72 return wrapper(f_outer) # -> T_Loader 

73 else: 

74 return wrapper # -> T_WrappedLoader 

75 

76 

77__all__ = ["register_loader", "LOADERS"]