[fix] Parse array is None if CFG conditionals when quering for memlets

dace/sdfg/sdfg.py

@@ -123,6 +123,7 @@ def memlets_in_ast(node: ast.AST, arrays: Dict[str, dt.Data], *, include_scalars
     """
     Generates a list of memlets from each of the subscripts that appear in the Python AST.
     Assumes the subscript slice can be coerced to a symbolic expression (e.g., no indirect access).
+    Can also parse a None check of the form `array is [not] None`.
 
     :param node: The AST node to find memlets in.
     :param arrays: A dictionary mapping array names to their data descriptors (a-la ``sdfg.arrays``)
@@ -133,10 +134,16 @@ def memlets_in_ast(node: ast.AST, arrays: Dict[str, dt.Data], *, include_scalars
 
     for subnode in ast.walk(node):
         if isinstance(subnode, ast.Subscript):
-            data = astutils.rname(subnode.value)
             data, slc = astutils.subscript_to_slice(subnode, arrays)
             subset = sbs.Range(slc)
             result.append(mm.Memlet(data=data, subset=subset))
+        elif (isinstance(subnode, ast.Compare) and len(subnode.ops) == 1
+              and isinstance(subnode.ops[0], (ast.Is, ast.IsNot)) and len(subnode.comparators) == 1
+              and isinstance(subnode.comparators[0], ast.Constant) and subnode.comparators[0].value is None):
+            # Parsing `array is [not] None`
+            data = astutils.rname(subnode.left)
+            if data in arrays:
+                result.append(mm.Memlet.from_array(data, arrays[data]))
         elif include_scalars and isinstance(subnode, ast.Name):
             data = astutils.rname(subnode)
             if data in arrays and isinstance(arrays[data], dace.data.Scalar):

tests/control_flow_test.py

@@ -1,4 +1,5 @@
 # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved.
+from typing import Optional
 import pytest
 import dace
 import numpy as np
@@ -339,6 +340,29 @@ def test_fsm():
         assert 'for ' not in code
 
 
+def test_optional_parameters():
+
+    def optional_parameters_func(A: dace.int32[3], B: Optional[dace.int32[3]] = None):
+        if B is None:
+            A[1] = 3
+        elif B is not None:
+            A[2] = 5
+
+    @dace.program
+    def optional_parameters_program(A: dace.int32[3], B: dace.int32[3]):
+        optional_parameters_func(A)
+        optional_parameters_func(A, B)
+
+    sdfg: dace.SDFG = optional_parameters_program.to_sdfg()
+    A = np.zeros((3, ), dtype=np.int32)
+    B = np.zeros((3, ), dtype=np.int32)
+
+    sdfg(A, B)
+    assert A[0] == 0
+    assert A[1] == 3
+    assert A[2] == 5
+
+
 if __name__ == '__main__':
     test_control_flow_basic()
     test_function_in_condition()
@@ -352,3 +376,4 @@ def test_fsm():
     test_ifchain_manual()
     #test_switchcase()
     test_fsm()
+    test_optional_parameters()