def __init__(self, start_date, end_date, user_config=None):
"""Initialize the generator."""
if not self.TEMPLATE:
raise AttributeError("Class attribute 'TEMPLATE' must be defined.")
if not self.TEMPLATE_KWARGS:
raise AttributeError(
"Class attribute 'TEMPLATE_KWARGS' must be defined.")
LOADER_LIST.append(PackageLoader("nise"))
if user_config:
LOADER_LIST.append(
FileSystemLoader(os.path.abspath(
os.path.dirname(user_config))))
env = Environment(loader=ChoiceLoader(LOADER_LIST))
env.globals["faker"] = faker_passthrough
default_template = env.get_template(self.TEMPLATE)
if user_config:
user_template = env.get_template(os.path.basename(user_config))
user_yaml = load_yaml(user_template.render(**self.TEMPLATE_KWARGS))
# sort lists of dicts so that generator class names align.
generators = user_yaml.get("generators")
user_yaml["generators"] = sorted(generators,
key=lambda d: list(d.keys()))
default_yaml = load_yaml(
default_template.render(**self.TEMPLATE_KWARGS))
config = deepupdate(
default_yaml, user_yaml
) # merge user-supplied static file with base template
else:
config = load_yaml(default_template.render(**self.TEMPLATE_KWARGS))
# handle special-cases in YAML config syntax
config = self._format_config(config)
# remove top-level class name
self.config = []
for generators in config.get("generators"):
for key, val in generators.items():
if key == type(self).__name__:
self.config.append(val)
self.start_date = self._set_date_config(start_date, "start")
self.end_date = self._set_date_config(end_date, "end")
self.hours = self._set_hours()
self.days = self._set_days()
LOG.debug("Current config: %s", pformat(self.config))
super().__init__()
def main():
"""Run data generation program."""
parser = create_parser()
args = parser.parse_args()
if args.log_level:
LOG.setLevel(LOG_VERBOSITY[args.log_level])
if not args.command:
parser.error('"yaml" or "report" argument must be specified')
elif args.command == "yaml":
yaml_main(args)
return
options = vars(args)
LOG.debug("Options are: %s", pformat(options))
if not (options.get("start_date") or options.get("static_report_file")):
parser.error("the following arguments are required: -s, --start-date")
_, provider_type = _validate_provider_inputs(parser, options)
run(provider_type, options)