Patching certain functions raises `NameError: name 'snowflake' is not defined`

[narciero]

Please answer these questions before submitting your issue. Thanks!

1. What version of Python are you using?

   Python 3.12.11

2. What operating system and processor architecture are you using?

   macOS-26.2-arm64-arm-64bit

3. What are the component versions in the environment (pip freeze)?

   snowflake-connector-python==4.3.0
   snowflake-ml-python==1.29.0
   snowflake-snowpark-python==1.47.0
   

4. What did you do?

   from snowflake.snowpark import Session
   from snowflake.snowpark.functions import col, square
   from snowflake.snowpark.mock import patch
   
   @patch(square)
   def mock_square(x):
      raise AssertionError("never reached")
   
   session = Session.builder.config("local_testing", True).getOrCreate()
   df = session.create_dataframe([[1]], schema=["x"])
   df.select(square(col("x"))).collect()  # NameError: name 'snowflake' is not defined

5. What did you expect to see?

   The patch function (mock_square) should be executed. Instead, the following error is thrown:

   NameError: name 'snowflake' is not defined
   

6. Can you set logging to DEBUG and collect the logs?

   N/A

Suspected Issue

• The ColumnOrName = Union["snowflake.snowpark.column.Column", str] typedef in _internal/type_utils.py
  uses a forward ref.
• When handle_function_expression calls typing.get_type_hints() on the
  original function being patched, it evaluates the forward ref in _functions/scalar_functions.py's globals
  where 'snowflake' is not a bare name.

[sfc-gh-sghosh]

Hello @narciero ,

Thanks for raising the issue.
I checked the code snippet and able to reproduce. Its actually working as expected.

With regular session: The output is [Row(SQUARE("X")=1.0)]
With local testing: The output is SnowparkLocalTestingException: Error executing mocked function 'square'. See error traceback for detailed information.

You are getting an error with local_testing because you are throwing an exception raise AssertionError("never reached")

Local Session (Expected Failure): The Local Testing framework use the @patch. Your function then does exactly what you told it to do: it crashes with an AssertionError. Whereas for a regular session, it ignores the @patch, and it uses the regular built-in square function.

So if you want you mock_square function same as a regular session, you can use the code below instead of throwing an error

from snowflake.snowpark import Session
from snowflake.snowpark.functions import col, square
from snowflake.snowpark.mock import patch

@patch(square)
def mock_square(x):
   #raise AssertionError("never reached")
   return x * x

session = Session.builder.config('local_testing', True).create()
df = session.create_dataframe([[1]], schema=["x"])
df.select(square(col("x"))).collect()  # NameError: name 'snowflake' is not defined

Output:
[Row(SQUARE("X")=1.0)]