class AssignmentExtractor(ast.NodeVisitor):
def __init__(self):
from pyfileconf.assignments.models.container import AssignmentStatementContainer
self.assigns = AssignmentStatementContainer([])
def visit_Assign(self, node):
from pyfileconf.assignments.models.statement import AssignmentStatement
self.assigns.append(AssignmentStatement.from_ast_assign(node))
def visit_AnnAssign(self, node):
from pyfileconf.assignments.models.statement import AssignmentStatement
self.assigns.append(AssignmentStatement.from_ast_assign(node))
def __init__(self,
d: dict = None,
name: str = None,
annotations: dict = None,
imports: ImportStatementContainer = None,
_file: ConfigFileBase = None,
begin_assignments: AssignmentStatementContainer = None,
klass: Optional[Type] = None,
always_import_strs: Optional[Sequence[str]] = None,
always_assign_strs: Optional[Sequence[str]] = None,
**kwargs):
if d is None:
d = {}
super().__init__(d, **kwargs)
if annotations is None:
annotations = {}
if imports is None:
imports = ImportStatementContainer([])
if begin_assignments is None:
begin_assignments = AssignmentStatementContainer([])
self.name = name
self.annotations = annotations
self.imports = imports
self._file = _file
self.begin_assignments = begin_assignments
self.klass = klass
self.always_import_strs = always_import_strs
self.always_assign_strs = always_assign_strs
self._applied_updates: Dict[str, Any] = {}
def save(self, config: 'ConfigBase'):
# Set existing assigns and body to empty because should be completely replaced by new specific class dict
file_str_obj = FileStr(config,
existing_assigns=AssignmentStatementContainer(
[]),
existing_imports=self.imports,
existing_body=[])
with open(self.filepath, 'w', newline='\n', encoding='utf8') as f:
f.write(file_str_obj.file_str)
def register(self):
# Store ast representation of file and file body
super().register()
# Store imports and assignments
if self._ast is not None:
self._imports = extract_imports_from_ast(self._ast)
self._assigns = extract_assignments_from_ast(self._ast)
else:
self._imports = ImportStatementContainer([])
self._assigns = AssignmentStatementContainer([])
def _add_always_imports_and_assigns_to_config(self, config: 'ConfigBase'):
"""
Note: inplace
"""
# Add always imports
[config.imports.add_if_missing(imp) for imp in self.always_imports]
# # Check if there are any extra assigns for items with this name
always_assigns = self.always_assigns.copy()
if self.name in self.always_assign_with_names_dict:
always_assigns.extend(
self.always_assign_with_names_dict[self.name])
# Add always assigns
# First handle begin assigns
begin_assigns = AssignmentStatementContainer(
[assign for assign in always_assigns if assign.prefer_beginning])
config.begin_assignments = begin_assigns
# Now handle the rest
# First get always assigns, annotations as dict
other_always_assigns = AssignmentStatementContainer([
assign for assign in always_assigns if not assign.prefer_beginning
])
always_defaults, always_annotations = other_always_assigns.to_default_dict_and_annotation_dict(
)
# Select assigns, annotations which are not already defined in config
new_defaults = {
key: value
for key, value in always_defaults.items() if key not in config
}
new_annotations = {
key: value
for key, value in always_annotations.items()
if key not in config.annotations
}
# Add to config
config.update(new_defaults)
config.annotations.update(new_annotations)
def _combine_imports_get_new_assignments(
self, import_assignment_obj,
existing_imports: ImportStatementContainer,
existing_assigns: AssignmentStatementContainer
) -> ImportsDoubleAssignsTuple:
all_imports = existing_imports.copy()
new_assigns_begin = AssignmentStatementContainer([])
new_assigns_end = AssignmentStatementContainer([])
possibly_new_imports, possibly_new_assigns = import_assignment_obj.as_imports_and_assignments(
)
# Checks to see whether should be added, and whether to beginning or end, then adds
[all_imports.add_if_missing(imp) for imp in possibly_new_imports]
for assign in possibly_new_assigns:
if not existing_assigns.contains_varname(assign.varname):
begin = getattr(assign, 'prefer_beginning', False)
if begin:
new_assigns_begin.append_if_missing(assign)
else:
new_assigns_end.append_if_missing(assign)
return all_imports, new_assigns_begin, new_assigns_end
def extract_import_statements_from_function_args_imports_and_assigns(
args: ast.arguments,
imports: ImportStatementContainer,
assigns: AssignmentStatementContainer,
current_module_section_path_str: str,
) -> ImportStatementContainer:
defaults_dict, annotation_dict = function_args_as_arg_and_annotation_dict(
args)
# Extract external names. These are the root names (not attrs, the base of attrs) for anything that is not builtin
external_names = _unique_external_names_from_default_dict_and_annotation_dict(
defaults_dict, annotation_dict)
out_imports = ImportStatementContainer([])
# External names may be either due to assignment or due to imports
# Handle names due to assignment
for name in external_names:
if assigns.contains_varname(name):
# External name is due to assignment in this module. Create import from this module
out_imports.append(
ObjectImportStatement.from_str(
f'from {current_module_section_path_str} import {name}'))
# Handle names which are imported into this file
for name in external_names:
import_or_none = imports.get_import_for_module_or_obj_name(name)
if import_or_none is not None:
out_imports.append(import_or_none)
# Sanity check, did we find all the imports?
n_external = len(external_names)
n_imports = len(out_imports)
if n_external != n_imports:
warnings.warn(
f'had {n_external} external names from function definition, only '
f'found {n_imports} imports. May be missing imports')
return out_imports
def as_imports_and_assignments(self) -> ImportsAndAssigns:
pipeline_dict_assign = AssignmentStatement.from_str(self['class_dict'])
assigns = AssignmentStatementContainer([pipeline_dict_assign])
return self.imports, assigns
def as_imports_and_assignments(self) -> ImportsAndAssigns:
assigns = AssignmentStatementContainer.from_dict_of_varnames_and_ast(
self, self.annotations)
return self.imports, self.begin_assignments + assigns
def __init__(self):
from pyfileconf.assignments.models.container import AssignmentStatementContainer
self.assigns = AssignmentStatementContainer([])