|
11 | 11 |
|
12 | 12 |
|
13 | 13 | class DataPluginConfig(base.Config): |
14 | | - data_dir = c.Type(str, default='data') |
| 14 | + sources = c.Type(dict, default={'data': 'data'}) |
15 | 15 |
|
16 | 16 |
|
17 | 17 | class DataPlugin(BasePlugin[DataPluginConfig]): |
18 | 18 | def __init__(self): |
19 | | - self.data = {} |
| 19 | + self.sources = {} |
20 | 20 | self.processors = {'.yml': yaml.safe_load, '.yaml': yaml.safe_load, '.json': json.load} |
21 | 21 |
|
22 | | - def set_data(self, keys, value): |
| 22 | + def load_sources(self): |
23 | 23 | """ |
24 | | - Set a value in the data attribute. |
| 24 | + Load all sources from the config file and load the data from the files. |
25 | 25 | """ |
26 | | - data = self.data |
27 | | - for key in keys[:-1]: |
28 | | - data = data.setdefault(key, {}) |
29 | | - data[keys[-1]] = value |
| 26 | + for source, path in self.config['sources'].items(): |
| 27 | + if not os.path.exists(path): |
| 28 | + log.warning(f"Mapping path '{path}' not found. Skipping.") |
| 29 | + elif os.path.isdir(path): |
| 30 | + self.load_folder(source, path) |
| 31 | + else: |
| 32 | + value = self.load_file(path) |
| 33 | + self.update_data(source, [], value) |
30 | 34 |
|
31 | | - def on_config(self, config: MkDocsConfig): |
32 | | - self.load_data(self.config.data_dir) |
33 | | - |
34 | | - macros_plugin = config.plugins.get('macros') |
35 | | - if macros_plugin: |
36 | | - macros_plugin.register_variables({'data': self.data}) |
| 35 | + def update_data(self, source: str, keys: list, value: any): |
| 36 | + """ |
| 37 | + Update the sources data with the given value. |
| 38 | + """ |
| 39 | + if len(keys) == 0: |
| 40 | + self.sources[source] = value |
37 | 41 | else: |
38 | | - log.warning( |
39 | | - "The macros plugin is not installed. The `data` variable won't be available in pages." |
40 | | - ) |
| 42 | + data = self.sources.setdefault(source, {}) |
| 43 | + for key in keys[:-1]: |
| 44 | + data = data.setdefault(key, {}) |
| 45 | + data[keys[-1]] = value |
41 | 46 |
|
42 | | - def load_data(self, path: str): |
| 47 | + def load_folder(self, source: str, path: str): |
43 | 48 | """ |
44 | 49 | Iterate over all files in the data directory |
45 | 50 | and load them into the data attribute. |
46 | 51 | """ |
47 | 52 | for root, _, files in os.walk(path): |
48 | 53 |
|
49 | 54 | keys = [] |
50 | | - if root != self.config.data_dir: |
51 | | - directory = os.path.relpath(root, self.config.data_dir) |
| 55 | + if root != path: |
| 56 | + directory = os.path.relpath(root, path) |
52 | 57 | keys = directory.split(os.sep) |
53 | 58 |
|
54 | 59 | for file in files: |
55 | 60 | value = self.load_file(os.path.join(root, file)) |
56 | 61 |
|
57 | 62 | filename, _ = os.path.splitext(file) |
58 | | - self.set_data(keys + [filename], value) |
| 63 | + self.update_data(source, keys + [filename], value) |
59 | 64 |
|
60 | 65 | def load_file(self, path: str): |
61 | 66 | """ |
62 | | - Load a file and return its content. |
| 67 | + Loads a file and processes it with the appropriate processor. |
63 | 68 | """ |
64 | 69 | _, extension = os.path.splitext(path) |
65 | 70 | with open(path, 'r') as file: |
66 | 71 | return self.processors[extension](file) |
67 | 72 |
|
| 73 | + def on_config(self, config: MkDocsConfig): |
| 74 | + self.load_sources() |
| 75 | + |
| 76 | + macros_plugin = config.plugins.get('macros') |
| 77 | + if macros_plugin: |
| 78 | + for source, data in self.sources.items(): |
| 79 | + macros_plugin.register_variables({source: data}) |
| 80 | + else: |
| 81 | + log.warning( |
| 82 | + "The macros plugin is not installed. The `data` variable won't be available in pages." |
| 83 | + ) |
| 84 | + |
68 | 85 | def on_page_context(self, context, page, config, nav): |
69 | | - context['data'] = self.data |
| 86 | + for source, data in self.sources.items(): |
| 87 | + context[source] = data |
70 | 88 | return context |
0 commit comments