Skip to content

Commit

Permalink
Fix python model parsing from meder (#9162)
Browse files Browse the repository at this point in the history
* fix: parsing f-strings in python models

* add unittest

* changelong

* pre-commit

---------

Co-authored-by: Meder Kamalov <[email protected]>
Co-authored-by: Kshitij Aranke <[email protected]>
  • Loading branch information
3 people authored Jan 31, 2024
1 parent dc3f608 commit 8717828
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 0 deletions.
6 changes: 6 additions & 0 deletions .changes/unreleased/Fixes-20231128-102111.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
kind: Fixes
body: Fix parsing f-strings in python models
time: 2023-11-28T10:21:11.596121-08:00
custom:
Author: mederka
Issue: "6976"
6 changes: 6 additions & 0 deletions core/dbt/parser/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,12 @@ def visit_Call(self, node: ast.Call) -> None:
for value in obj.values:
if isinstance(value, ast.Call):
self.visit_Call(value)
# support dbt function calls in f-strings
elif isinstance(obj, ast.JoinedStr):
for value in obj.values:
if isinstance(value, ast.FormattedValue) and isinstance(value.value, ast.Call):
self.visit_Call(value.value)

# visit node.func.value if we are at an call attr
if isinstance(node.func, ast.Attribute):
self.attribute_helper(node.func)
Expand Down
21 changes: 21 additions & 0 deletions tests/unit/test_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -798,6 +798,7 @@ def model(dbt, session):
a_dict = {'test2': dbt.ref('test2')}
df5 = {'test2': dbt.ref('test3')}
df6 = [dbt.ref("test4")]
f"{dbt.ref('test5')}"
df = df0.limit(2)
return df
Expand Down Expand Up @@ -861,6 +862,17 @@ def model(dbt, session):
return dbt.ref("some_model"), dbt.ref("some_other_model")
"""

python_model_f_string = """
# my_python_model.py
import pandas as pd
def model(dbt, fal):
dbt.config(materialized="table")
print(f"my var: {dbt.config.get('my_var')}") # Prints "my var: None"
df: pd.DataFrame = dbt.ref("some_model")
return df
"""

python_model_no_return = """
def model(dbt, session):
dbt.config(materialized='table')
Expand Down Expand Up @@ -985,6 +997,7 @@ def test_python_model_parse(self):
RefArgs("test2"),
RefArgs("test3"),
RefArgs("test4"),
RefArgs("test5"),
],
sources=[["test", "table1"]],
)
Expand All @@ -1001,6 +1014,14 @@ def test_python_model_config(self):
node = list(self.parser.manifest.nodes.values())[0]
self.assertEqual(node.config.to_dict()["config_keys_used"], ["param_1", "param_2"])

def test_python_model_f_string_config(self):
block = self.file_block_for(python_model_f_string, "nested/py_model.py")
self.parser.manifest.files[block.file.file_id] = block.file

self.parser.parse_file(block)
node = list(self.parser.manifest.nodes.values())[0]
self.assertEqual(node.config.to_dict()["config_keys_used"], ["my_var"])

def test_python_model_config_with_defaults(self):
block = self.file_block_for(python_model_config_with_defaults, "nested/py_model.py")
self.parser.manifest.files[block.file.file_id] = block.file
Expand Down

0 comments on commit 8717828

Please sign in to comment.