Skip to content

Commit

Permalink
Add types_to_file_type + tweak annotation handling
Browse files Browse the repository at this point in the history
  • Loading branch information
Cyrilvallez committed Oct 29, 2024
1 parent 5f060ae commit c72e112
Showing 1 changed file with 10 additions and 3 deletions.
13 changes: 10 additions & 3 deletions utils/modular_model_converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -490,9 +490,13 @@ def __init__(self, class_name: str, global_names: set | None):
def visit_Name(self, node):
if node.value != self.class_name and node.value in self.global_names:
parent_node = self.get_metadata(cst.metadata.ParentNodeProvider, node)
# If it is only an annotation, do not add dependency
if not m.matches(parent_node, m.Annotation()):
self.dependencies.add(node.value)
# If it is only an annotation inside a method definition, do not add dependency (however, do it for
# annotations that are variable definitions, i.e. for Kwargs classes)
if m.matches(parent_node, m.Annotation()):
grand_parent = self.get_metadata(cst.metadata.ParentNodeProvider, parent_node)
if m.matches(grand_parent, m.Param() | m.FunctionDef()):
return
self.dependencies.add(node.value)


def dependencies_for_class_node(node: cst.ClassDef, global_names: set) -> set:
Expand Down Expand Up @@ -890,6 +894,9 @@ def replace_class_node(mapper: ModelFileMapper, class_node: cst.ClassDef, rename
"Processor": "processing",
"ImageProcessor": "image_processing",
"FeatureExtractor": "feature_extractor",
"ProcessorKwargs": "processing",
"ImagesKwargs": "processing",
"TextKwargs": "processing",
}


Expand Down

0 comments on commit c72e112

Please sign in to comment.