import logging import pathlib import json import yaml from yaml.loader import SafeLoader from sys import exit from pelican.generators import Generator from .file_formats import JSON, YAML log = logging.getLogger(__name__) log.setLevel(logging.INFO) class DataGenerator(Generator): """ Load data from files """ SUPPORTED_FORMATS = [JSON, YAML] CONTEXT_PREFIX = "DATA_" def __init__( self, context, settings, path, theme, output_path, readers_cache_name="", **kwargs, ): super().__init__( context, settings, path, theme, output_path, readers_cache_name=readers_cache_name, **kwargs, ) log.info("PLUGIN: pelican-data-files was successfully loaded") self.settings.setdefault("DATA_FILES_DIR", "data") def _is_valid_file(self, file): """Checks if file format is supported. params: - file -- pathlib.Path object """ for file_format in self.SUPPORTED_FORMATS: if file.suffix in file_format["extensions"]: return True return False def _get_data_files(self): """Return list of valid files to load into context""" data_dir = pathlib.Path(self.settings["DATA_FILES_DIR"]) valid_files = [] # turn path into absolute if not already if not data_dir.is_absolute(): data_dir = pathlib.Path(self.settings["PATH"]).joinpath(data_dir) # check if path exists if not data_dir.exists(): log.error("pelican-data-files: DATA_FILES_DIR path doesn't exists.") exit(1) if not data_dir.is_dir(): log.error("pelican-data-files: DATA_FILES_DIR path isn't a directory.") exit(1) # return all valid files in path # TODO check for duplicates (eg: profile.json and profile.yaml) for file in data_dir.iterdir(): if self._is_valid_file(file): valid_files.append(file) return valid_files def _format_filename(self, file): """Format context var name from filename. params: - file -- pathlib.Path object """ return file.stem.replace(".", "_").upper() def _read_file(self, file): """Read and parse data from file. params: - file -- pathlib.Path object """ with file.open() as f: for file_format in self.SUPPORTED_FORMATS: if file_format["name"] == "json" and file.suffix in file_format["extensions"]: try: return json.load(f) except ValueError: return None elif file_format["name"] == "yaml" and file.suffix in file_format["extensions"]: try: return yaml.load(f, Loader=SafeLoader) except ValueError: return None def _add_data_to_context(self, name, data): """Add data into context. params: - name -- str - data -- dict """ ctx_name = self.CONTEXT_PREFIX + name self.context[ctx_name] = data def generate_context(self): """Generate context from data files""" for file in self._get_data_files(): name = self._format_filename(file) data = self._read_file(file) if data: self._add_data_to_context(name, data) log.info(f"{file.name} was loaded.") else: log.error(f"{file.name} wasn't loaded.") def get_generators(pelican_object): return DataGenerator