Support custom dosctrings in modular (#36726)

* Override docstrings in modular if not none

* Update doc
This commit is contained in:
Yoni Gozlan
2025-03-18 14:00:54 -04:00
committed by GitHub
parent 00915d3041
commit 12f2ebef63
5 changed files with 27 additions and 25 deletions

View File

@@ -537,6 +537,9 @@ def find_all_dependencies(
# Top-level variables that match the following patterns will always use the value in the `modular_xxx.py` file
ASSIGNMENTS_REGEX_TO_KEEP = [r"_CHECKPOINT", r"_EXPECTED", r"_FOR_DOC"]
# Top-level variables that match the following patterns will use the value in the `modular_xxx.py` file only if they are not None
ASSIGNMENTS_REGEX_TO_KEEP_IF_NOT_NONE = [r"_DOCSTRING"]
class ClassDependencyMapper(CSTVisitor):
"""A visitor which is designed to analyze a single class node to get all its dependencies that are shared with the set of
@@ -854,13 +857,17 @@ class ModelFileMapper(ModuleMapper):
"""Update the global nodes with the assignment from the modular file.
Merging rule: if any assignment with the same name was redefined in the modular, we use it and its dependencies ONLY if it matches
a pattern in `ASSIGNMENTS_REGEX_TO_KEEP`. Otherwise, we use the original value and dependencies. This rule was chosen to avoid having to rewrite the
big docstrings.
a pattern in `ASSIGNMENTS_REGEX_TO_KEEP_IF_NOT_NONE` and its value is not None, or if it matches a pattern in `ASSIGNMENTS_REGEX_TO_KEEP.
Otherwise, we use the original value and dependencies. This rule was chosen to avoid having to rewrite the big docstrings.
"""
for assignment, node in assignments.items():
should_keep = any(re.search(pattern, assignment) for pattern in ASSIGNMENTS_REGEX_TO_KEEP)
if should_keep or assignment not in self.assignments:
should_keep_if_not_none = any(
re.search(pattern, assignment) for pattern in ASSIGNMENTS_REGEX_TO_KEEP_IF_NOT_NONE
) and not (hasattr(node.body[0].value, "value") and node.body[0].value.value == "None")
if should_keep or should_keep_if_not_none or assignment not in self.assignments:
self.assignments[assignment] = node
if assignment in object_mapping:
self.object_dependency_mapping[assignment] = object_mapping[assignment]