def write_output(
path: str, stat_set: BaseModel, is_local: bool = False
) -> int:
if settings.export_local:
is_local = True
if hasattr(stat_set, "json"):
indent = None
if settings.debug:
indent = 4
write_content = stat_set.json(exclude_unset=True, indent=indent)
else:
write_content = json.dumps(stat_set)
byte_count = 0
if is_local:
byte_count = write_to_local(path, write_content)
else:
byte_count = write_to_s3(path, write_content)
return byte_count
def _add_parameter(self, sess: Session, task: Task, param: BaseModel):
if task is None:
raise PPLException(ErrorCode.TASK_NOT_EXISTS, None)
data = Data()
data.data = param.json()
sess.add(data)
if task.parameter is not None:
db_param = task.parameter
else:
db_param = Parameter()
db_param.data = data
db_param.task = task
sess.add(db_param)
def _dump_result(self, sess: Session, task_id: int, result: BaseModel):
task = self._get_task(sess, task_id)
if task is None:
raise PPLException(ErrorCode.TASK_NOT_EXISTS, task_id)
data = Data()
data.data = result.json()
sess.add(data)
if task.result is not None:
db_result = task.result
else:
db_result = Result()
db_result.task = task
db_result.data = data
sess.add(db_result)
def _save(service_name: str, filename: Path, model: BaseModel):
output_path = filename
if has_multiple_services:
output_path = filename.parent / service_name / filename.name
output_path.parent.mkdir(parents=True, exist_ok=True)
click.echo(f"Creating {output_path} ...", nl=False)
with output_path.open("wt") as fh:
data = json.loads(
model.json(exclude_unset=True,
by_alias=True,
exclude_none=True))
yaml.safe_dump(data, fh, sort_keys=False)
click.echo("DONE")
def write_output(
path: str,
stat_set: BaseModel,
is_local: bool = False,
exclude_unset: bool = True,
exclude: set = None,
) -> int:
if settings.export_local:
is_local = True
if hasattr(stat_set, "json"):
indent = None
if settings.debug:
indent = 4
write_content = stat_set.json(exclude_unset=exclude_unset,
indent=indent,
exclude=exclude)
else:
write_content = json.dumps(stat_set)
byte_count = 0
if is_local:
byte_count = write_to_local(path, write_content)
elif isinstance(stat_set, str):
byte_count = write_to_s3(stat_set, path)
elif isinstance(stat_set, OpennemDataSet):
byte_count = write_statset_to_s3(stat_set,
path,
exclude_unset=exclude_unset,
exclude=exclude)
elif isinstance(stat_set, BaseModel):
byte_count = write_to_s3(write_content, path)
else:
raise Exception(
"Do not know how to write content of this type to output")
return byte_count
def save_data(source: str, city: str, data: BaseModel) -> Path:
file_path = make_filepath(source=source, city=city)
with open(file_path, "w") as json_file:
json.dump(data.json(), json_file)
return file_path
def serialize_payload(self, payload: BaseModel):
serialized = payload.json()
return serialized.encode("utf-8")
def publish(self, payload: BaseModel, topic_name: str):
topic_path = self.get_topic_path(topic_name)
data = payload.json()
Logger.info(f"Published virtual pub/sub message to '{topic_path}")
Logger.info(data)
def encode(model: BaseModel) -> str:
"""Convert the given model to a string"""
return model.json()