def visit_Module(self, node: libcst.Module) -> None:
# Do a preliminary pass to gather the imports we already have
gatherer = GatherImportsVisitor(self.context)
node.visit(gatherer)
self.all_imports = gatherer.all_imports
self.module_imports = self.module_imports - gatherer.module_imports
for module, alias in gatherer.module_aliases.items():
if module in self.module_aliases and self.module_aliases[
module] == alias:
del self.module_aliases[module]
for module, aliases in gatherer.alias_mapping.items():
for (obj, alias) in aliases:
if (module in self.alias_mapping
and (obj, alias) in self.alias_mapping[module]):
self.alias_mapping[module].remove((obj, alias))
if len(self.alias_mapping[module]) == 0:
del self.alias_mapping[module]
for module, imports in gatherer.object_mapping.items():
if module not in self.module_mapping:
# We don't care about this import at all
continue
elif "*" in imports:
# We already implicitly are importing everything
del self.module_mapping[module]
else:
# Lets figure out what's left to import
self.module_mapping[
module] = self.module_mapping[module] - imports
if not self.module_mapping[module]:
# There's nothing left, so lets delete this work item
del self.module_mapping[module]
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
"""
Collect type annotations from all stubs and apply them to ``tree``.
Gather existing imports from ``tree`` so that we don't add duplicate imports.
"""
import_gatherer = GatherImportsVisitor(CodemodContext())
tree.visit(import_gatherer)
existing_import_names = _get_import_names(import_gatherer.all_imports)
context_contents = self.context.scratch.get(
ApplyTypeAnnotationsVisitor.CONTEXT_KEY
)
if context_contents:
stub, overwrite_existing_annotations = context_contents
self.overwrite_existing_annotations = (
self.overwrite_existing_annotations or overwrite_existing_annotations
)
visitor = TypeCollector(existing_import_names, self.context)
stub.visit(visitor)
self.annotations.function_annotations.update(visitor.function_annotations)
self.annotations.attribute_annotations.update(visitor.attribute_annotations)
self.annotations.class_definitions.update(visitor.class_definitions)
tree_with_imports = AddImportsVisitor(self.context).transform_module(tree)
return tree_with_imports.visit(self)
def _annotate_source(stubs: cst.Module, source: cst.Module) -> cst.Module:
visitor = TypeCollector()
stubs.visit(visitor)
transformer = TypeTransformer(
visitor.function_annotations, visitor.attribute_annotations, visitor.imports
)
return source.visit(transformer)
def visit_Module(self, node: cst.Module) -> None:
# Collect current list of imports
gatherer = GatherImportsVisitor(self.context)
node.visit(gatherer)
# Store list of symbols imported from wx package
self.wx_imports = gatherer.object_mapping.get("wx", set())
def visit_Module(self, node: cst.Module) -> bool:
export_collector = GatherExportsVisitor(self.context)
node.visit(export_collector)
self._exported_names = export_collector.explicit_exported_objects
annotation_visitor = GatherNamesFromStringAnnotationsVisitor(
self.context, typing_functions=self._typing_functions)
node.visit(annotation_visitor)
self._string_annotation_names = annotation_visitor.names
return True
def visit_Module(self, node: cst.Module) -> bool:
visitor = FullyQualifiedNameVisitor(self, self.module_name)
node.visit(visitor)
self.set_metadata(
node,
{
QualifiedName(name=self.module_name,
source=QualifiedNameSource.LOCAL)
},
)
return True
def _annotate_source(stubs: cst.Module, source: cst.Module) -> cst.Module:
import_visitor = ImportCollector()
source.visit(import_visitor)
visitor = TypeCollector(import_visitor.existing_imports)
stubs.visit(visitor)
transformer = TypeTransformer(
visitor.function_annotations,
visitor.attribute_annotations,
visitor.imports,
visitor.class_definitions,
)
return source.visit(transformer)
def transform_module_impl(
self,
tree: cst.Module,
) -> cst.Module:
"""
Collect type annotations from all stubs and apply them to ``tree``.
Gather existing imports from ``tree`` so that we don't add duplicate imports.
"""
import_gatherer = GatherImportsVisitor(CodemodContext())
tree.visit(import_gatherer)
existing_import_names = _get_imported_names(
import_gatherer.all_imports)
context_contents = self.context.scratch.get(
ApplyTypeAnnotationsVisitor.CONTEXT_KEY)
if context_contents is not None:
(
stub,
overwrite_existing_annotations,
use_future_annotations,
strict_posargs_matching,
strict_annotation_matching,
) = context_contents
self.overwrite_existing_annotations = (
self.overwrite_existing_annotations
or overwrite_existing_annotations)
self.use_future_annotations = (self.use_future_annotations
or use_future_annotations)
self.strict_posargs_matching = (self.strict_posargs_matching
and strict_posargs_matching)
self.strict_annotation_matching = (self.strict_annotation_matching
or strict_annotation_matching)
visitor = TypeCollector(existing_import_names, self.context)
cst.MetadataWrapper(stub).visit(visitor)
self.annotations.update(visitor.annotations)
if self.use_future_annotations:
AddImportsVisitor.add_needed_import(self.context, "__future__",
"annotations")
tree_with_imports = AddImportsVisitor(
self.context).transform_module(tree)
tree_with_changes = tree_with_imports.visit(self)
# don't modify the imports if we didn't actually add any type information
if self.annotation_counts.any_changes_applied():
return tree_with_changes
else:
return tree
def visit_Module(self, node: cst.Module) -> Optional[bool]:
node.visit(ExpressionContextVisitor(self, ExpressionContext.LOAD))
def visit_Module(self, node: cst.Module) -> Optional[bool]:
node.visit(ParentNodeVisitor(self))
def transform_module_impl(self, tree: libcst.Module) -> libcst.Module:
return tree.visit(self)
def _transform_module_impl(self, tree: cst.Module) -> cst.Module:
return tree.visit(self)
def visit_Module(self, node: cst.Module) -> None:
visitor = GatherUnusedImportsVisitor(self.context)
node.visit(visitor)
self._unused_imports = {k: v for (k, v) in visitor.unused_imports}
def visit_Module(self, node: cst.Module) -> None:
object_visitor = GatherExportsVisitor(self.context)
node.visit(object_visitor)
self.exported_objects = object_visitor.explicit_exported_objects
def visit_Module(self, node: Module) -> bool:
comment_visitor = GatherCommentsVisitor(
self.context, DEFAULT_SUPPRESS_COMMENT_REGEX)
node.visit(comment_visitor)
self._ignored_lines = set(comment_visitor.comments.keys())
return True
def visit_Module(self, node: cst.Module) -> Optional[bool]:
visitor = ScopeVisitor(self)
node.visit(visitor)
visitor.infer_accesses()
def visit_Module(self, node: cst.Module) -> Optional[bool]:
visitor = QualifiedNameVisitor(self)
node.visit(visitor)
def _get_attributes(source: cst.Module) -> Dict[str, List[str]]:
visitor = NonFinalAttributeCollector()
source.visit(visitor)
return visitor.attributes
