Python: small test of local flow

python/ql/src/semmle/code/python/dataflow/DataFlow.qll

@@ -11,8 +11,8 @@
  *
  * To use global (interprocedural) data flow, extend the class
  * `DataFlow::Configuration` as documented on that class. To use local
- * (intraprocedural) data flow, call `DataFlow::localFlowStep*` with
- *  arguments of type `DataFlow::Node`.
+ * (intraprocedural) data flow, call `DataFlow::localFlow` or
+ * `DataFlow::localFlowStep` with arguments of type `DataFlow::Node`.
  */
 
 import python

python/ql/src/semmle/code/python/dataflow/DataFlow2.qll

@@ -11,8 +11,8 @@
  *
  * To use global (interprocedural) data flow, extend the class
  * `DataFlow::Configuration` as documented on that class. To use local
- * (intraprocedural) data flow, call `DataFlow::localFlowStep*` with
- *  arguments of type `DataFlow::Node`.
+ * (intraprocedural) data flow, call `DataFlow::localFlow` or
+ * `DataFlow::localFlowStep` with arguments of type `DataFlow::Node`.
  */
 
 import python

python/ql/src/semmle/code/python/dataflow/internal/DataFlowImplSpecific.qll

@@ -8,4 +8,5 @@ module Private {
 
 module Public {
   import DataFlowPublic
+  import DataFlowUtil
 }

python/ql/src/semmle/code/python/dataflow/internal/DataFlowPrivate.qll

@@ -42,13 +42,15 @@ abstract class PostUpdateNode extends Node {
   abstract Node getPreUpdateNode();
 }
 
+private newtype TReturnKind = TNormalReturnKind()
+
 /**
  * A return kind. A return kind describes how a value can be returned
- * from a callable.
+ * from a callable. For Python, this is simply a method return.
  */
-abstract class ReturnKind extends string {
-  /** Gets a textual representation of this position. */
-  ReturnKind() { this = "ReturnKind" }
+class ReturnKind extends TReturnKind {
+  /** Gets a textual representation of this return kind. */
+  string toString() { result = "return" }
 }
 
 /** A data flow node that represents a value returned by a callable. */

python/ql/src/semmle/code/python/dataflow/internal/DataFlowUtil.qll

@@ -0,0 +1,18 @@
+/**
+ * Contains utility functions for writing data flow queries
+ */
+
+import DataFlowPrivate
+import DataFlowPublic
+
+/**
+ * Holds if data flows from `nodeFrom` to `nodeTo` in exactly one local
+ * (intra-procedural) step.
+ */
+predicate localFlowStep(Node nodeFrom, Node nodeTo) { simpleLocalFlowStep(nodeFrom, nodeTo) }
+
+/**
+ * Holds if data flows from `source` to `sink` in zero or more local
+ * (intra-procedural) steps.
+ */
+predicate localFlow(Node source, Node sink) { localFlowStep*(source, sink) }

python/ql/test/library-tests/dataflow/local.ql

@@ -0,0 +1,10 @@
+import python
+import semmle.code.python.dataflow.DataFlow
+
+from
+  DataFlow::Node fromNode,
+  DataFlow::Node toNode
+where
+  DataFlow::localFlowStep(fromNode, toNode)
+select
+  fromNode, toNode

python/ql/test/library-tests/dataflow/test.py

@@ -0,0 +1,8 @@
+a = 3
+b = a
+
+def f(x):
+  y = x + 2 # would expect flow to here from x
+  return y - 2 # would expect flow to here from y
+
+c = f(a)