generators
trestle.core.generators
¤
Capabilities to allow the generation of various oscal objects.
TG
¤
logger
¤
sample_base64
¤
sample_base64_value
¤
sample_date_value
¤
sample_method
¤
sample_observation_type_valid_value
¤
sample_task_valid_value
¤
type_base64
¤
Functions¤
generate_sample_model(model, include_optional=False, depth=-1)
¤
Given a model class, generate an object of that class with sample values.
Can generate optional variables with an enabled flag. Any array objects will have a single entry injected into it.
Note: Trestle generate will not activate recursive loops irrespective of the depth flag.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model |
Union[Type[~TG], List[~TG], Dict[str, ~TG]] |
The model type provided. Typically for a user as an OscalBaseModel Subclass. |
required |
include_optional |
bool |
Whether or not to generate optional fields. |
False |
depth |
int |
Depth of the tree at which optional fields are generated. Negative values (default) removes the limit. |
-1 |
Returns:
Type | Description |
---|---|
~TG |
The generated instance with a pro-forma values filled out as best as possible. |
Source code in trestle/core/generators.py
def generate_sample_model(
model: Union[Type[TG], List[TG], Dict[str, TG]], include_optional: bool = False, depth: int = -1
) -> TG:
"""Given a model class, generate an object of that class with sample values.
Can generate optional variables with an enabled flag. Any array objects will have a single entry injected into it.
Note: Trestle generate will not activate recursive loops irrespective of the depth flag.
Args:
model: The model type provided. Typically for a user as an OscalBaseModel Subclass.
include_optional: Whether or not to generate optional fields.
depth: Depth of the tree at which optional fields are generated. Negative values (default) removes the limit.
Returns:
The generated instance with a pro-forma values filled out as best as possible.
"""
effective_optional = include_optional and not depth == 0
model_type = model
# This block normalizes model type down to
if utils.is_collection_field_type(model): # type: ignore
model_type = utils.get_origin(model) # type: ignore
model = utils.get_inner_type(model) # type: ignore
model = cast(TG, model) # type: ignore
model_dict = {} # type: ignore
# this block is needed to avoid situations where an inbuilt is inside a list / dict.
# the only time dict ever appears is with include_all, which is handled specially
# the only type of collection possible after OSCAL 1.0.0 is list
if safe_is_sub(model, OscalBaseModel):
for field in model.__fields__: # type: ignore
if model_type in [OscalVersion]:
model_dict[field] = OSCAL_VERSION
break
if field == 'include_all':
if include_optional:
model_dict[field] = {}
continue
outer_type = model.__fields__[field].outer_type_ # type: ignore
# next appears to be needed for python 3.7
if utils.get_origin(outer_type) == Union:
outer_type = outer_type.__args__[0]
if model.__fields__[field].required or effective_optional: # type: ignore
# FIXME could be ForwardRef('SystemComponentStatus')
if utils.is_collection_field_type(outer_type):
inner_type = utils.get_inner_type(outer_type)
if inner_type == model:
continue
model_dict[field] = generate_sample_model(
outer_type, include_optional=include_optional, depth=depth - 1
)
elif is_by_type(outer_type):
model_dict[field] = generate_sample_value_by_type(outer_type, field)
elif safe_is_sub(outer_type, OscalBaseModel):
model_dict[field] = generate_sample_model(
outer_type, include_optional=include_optional, depth=depth - 1
)
else:
# Handle special cases (hacking)
if model_type in [Base64Datatype]:
model_dict[field] = sample_base64_value
elif model_type in [Base64]:
if field == 'filename':
model_dict[field] = sample_base64.filename
elif field == 'media_type':
model_dict[field] = sample_base64.media_type
elif field == 'value':
model_dict[field] = sample_base64.value
elif model_type in [DateDatatype]:
model_dict[field] = sample_date_value
# Hacking here:
# Root models should ideally not exist, however, sometimes we are stuck with them.
# If that is the case we need sufficient information on the type in order to generate a model.
# E.g. we need the type of the container.
elif field == '__root__' and hasattr(model, '__name__'):
model_dict[field] = generate_sample_value_by_type(
outer_type, str_utils.classname_to_alias(model.__name__, AliasMode.FIELD)
)
else:
model_dict[field] = generate_sample_value_by_type(outer_type, field)
# Note: this assumes list constrains in oscal are always 1 as a minimum size. if two this may still fail.
else:
if model_type is list:
return [generate_sample_value_by_type(model, '')] # type: ignore
if model_type is dict:
return {const.REPLACE_ME: generate_sample_value_by_type(model, '')} # type: ignore
raise err.TrestleError('Unhandled collection type.')
if model_type is list:
return [model(**model_dict)] # type: ignore
if model_type is dict:
return {const.REPLACE_ME: model(**model_dict)} # type: ignore
return model(**model_dict) # type: ignore
generate_sample_value_by_type(type_, field_name)
¤
Given a type, return sample value.
Includes the Optional use of passing down a parent_model
Source code in trestle/core/generators.py
def generate_sample_value_by_type(
type_: type,
field_name: str,
) -> Union[datetime, bool, int, str, float, Enum]:
"""Given a type, return sample value.
Includes the Optional use of passing down a parent_model
"""
# FIXME: Should be in separate generator module as it inherits EVERYTHING
if is_enum_method(type_):
return sample_method
if is_enum_task_valid_value(type_):
return sample_task_valid_value
if is_enum_observation_type_valid_value(type_):
return sample_observation_type_valid_value
if type_ is Base64:
return sample_base64
if type_ is datetime:
return datetime.now().astimezone()
if type_ is bool:
return False
if type_ is int:
return 0
if type_ is float:
return 0.00
if safe_is_sub(type_, ConstrainedStr) or (hasattr(type_, '__name__') and 'ConstrainedStr' in type_.__name__):
# This code here is messy. we need to meet a set of constraints. If we do
# TODO: handle regex directly
if 'uuid' == field_name:
return str(uuid.uuid4())
# some things like location_uuid in lists arrive here with field_name=''
if type_.regex and type_.regex.pattern.startswith('^[0-9A-Fa-f]{8}'): # type: ignore
return const.SAMPLE_UUID_STR
if field_name == 'date_authorized':
return str(date.today().isoformat())
if field_name == 'oscal_version':
return OSCAL_VERSION
if 'uuid' in field_name:
return const.SAMPLE_UUID_STR
# Only case where are UUID is required but not in name.
if field_name.rstrip('s') == 'member_of_organization':
return const.SAMPLE_UUID_STR
return const.REPLACE_ME
if hasattr(type_, '__name__') and 'ConstrainedIntValue' in type_.__name__:
# create an int value as close to the floor as possible does not test upper bound
multiple = type_.multiple_of if type_.multiple_of else 1 # type: ignore # default to every integer
# this command is a bit of a problem
floor = type_.ge if type_.ge else 0 # type: ignore
floor = type_.gt + 1 if type_.gt else floor # type: ignore
if math.remainder(floor, multiple) == 0:
return floor
return (floor + 1) * multiple
if safe_is_sub(type_, Enum):
# keys and values diverge due to hypens in oscal names
return type_(list(type_.__members__.values())[0]) # type: ignore
if type_ is str:
if field_name == 'oscal_version':
return OSCAL_VERSION
return const.REPLACE_ME
if type_ is pydantic.v1.networks.EmailStr:
return pydantic.v1.networks.EmailStr('dummy@sample.com')
if type_ is pydantic.v1.networks.AnyUrl:
# TODO: Cleanup: this should be usable from a url.. but it's not inuitive.
return pydantic.v1.networks.AnyUrl('https://sample.com/replaceme.html', scheme='http', host='sample.com')
if type_ is list:
raise err.TrestleError(f'Unable to generate sample for type {type_}')
# default to empty dict for anything else
return {} # type: ignore
is_by_type(model_type)
¤
Check for by type.
Source code in trestle/core/generators.py
def is_by_type(model_type: Union[Type[TG], List[TG], Dict[str, TG]]) -> bool:
"""Check for by type."""
rval = False
if model_type == type_base64:
rval = True
return rval
is_enum_method(type_)
¤
Test for method.
Source code in trestle/core/generators.py
def is_enum_method(type_: type) -> bool:
"""Test for method."""
rval = False
if utils.get_origin(type_) == Union:
args = typing.get_args(type_)
for arg in args:
if "<enum 'Methods'>" == f'{arg}':
rval = True
break
return rval
is_enum_observation_type_valid_value(type_)
¤
Test for observation type valid value.
Source code in trestle/core/generators.py
def is_enum_observation_type_valid_value(type_: type) -> bool:
"""Test for observation type valid value."""
rval = False
if utils.get_origin(type_) == Union:
args = typing.get_args(type_)
for arg in args:
if "<enum 'ObservationTypeValidValues'>" == f'{arg}':
rval = True
break
return rval
is_enum_task_valid_value(type_)
¤
Test for task valid value.
Source code in trestle/core/generators.py
def is_enum_task_valid_value(type_: type) -> bool:
"""Test for task valid value."""
rval = False
if utils.get_origin(type_) == Union:
args = typing.get_args(type_)
for arg in args:
if "<enum 'TaskValidValues'>" == f'{arg}':
rval = True
break
return rval
safe_is_sub(sub, parent)
¤
Is this a subclass of parent.
Source code in trestle/core/generators.py
def safe_is_sub(sub: Any, parent: Any) -> bool:
"""Is this a subclass of parent."""
is_class = inspect.isclass(sub)
return is_class and issubclass(sub, parent)
handler: python