You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
139 lines
3.7 KiB
Python
139 lines
3.7 KiB
Python
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
|