Skip to content

Commit

Permalink
fix: improve dataclass attribute handling in RestrictedPython transfo…
Browse files Browse the repository at this point in the history
…rmer

Co-Authored-By: Aaron <AJ> Steers <aj@airbyte.io>
  • Loading branch information
devin-ai-integration[bot] and aaronsteers committed Jan 22, 2025
1 parent f7afb52 commit 672dcd3
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions airbyte_cdk/sources/declarative/parsers/custom_code_compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,6 @@ def check_name(
self,
node: ast.AST,
name: str,
allow_magic_methods: bool = True,
*args: Any,
**kwargs: Any,
) -> ast.AST:
Expand All @@ -132,7 +131,6 @@ def check_name(
Args:
node: The AST node being checked
name: The name being validated
allow_magic_methods: Whether to allow magic methods (defaults to True)
*args: Additional positional arguments
**kwargs: Additional keyword arguments
"""
Expand Down Expand Up @@ -249,7 +247,9 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
allowed_nodes.append(self.visit_AnnAssign(n))
# Allow function definitions (like __post_init__)
elif isinstance(n, ast.FunctionDef):
allowed_nodes.append(self.visit(n))
# For function definitions, we need to allow attribute access
n.body = [self.visit(stmt) for stmt in n.body]
allowed_nodes.append(n)
# Allow docstrings
elif isinstance(n, ast.Expr) and isinstance(n.value, ast.Str):
allowed_nodes.append(n)
Expand All @@ -265,7 +265,11 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
)
allowed_nodes.append(self.visit_AnnAssign(ann_assign))
else:
# Allow attribute assignments within dataclasses
allowed_nodes.append(self.visit(n))
# Allow attribute statements within dataclasses
elif isinstance(n, ast.Attribute):
allowed_nodes.append(self.visit_Attribute(n))
else:
allowed_nodes.append(self.visit(n))

Expand Down

0 comments on commit 672dcd3

Please sign in to comment.