diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll b/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll
index 2f98ab70719..95434b05451 100644
--- a/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll
+++ b/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll
@@ -202,11 +202,18 @@ module TypeTrackingInput implements Shared::TypeTrackingInput<Location> {
    */
   predicate returnStep(Node nodeFrom, LocalSourceNode nodeTo) {
     exists(DataFlowPrivate::ExtractedDataFlowCall call |
-      nodeFrom.(DataFlowPrivate::ReturnNode).getEnclosingCallable() = call.getCallable() and
+      returnNodeEnclosingCallable(nodeFrom) = call.getCallable() and
       nodeTo.(DataFlowPublic::CfgNode).getNode() = call.getNode()
     )
   }
 
+  pragma[nomagic]
+  private DataFlowDispatch::DataFlowCallable returnNodeEnclosingCallable(
+    DataFlowPrivate::ReturnNode returnNode
+  ) {
+    result = returnNode.getEnclosingCallable()
+  }
+
   /**
    * Holds if `nodeFrom` is being written to the `content` content of the object in `nodeTo`.
    */