Python: Add test of type-tracking self in methods

python/ql/test/experimental/dataflow/typetracking/test.py

@@ -69,3 +69,35 @@ def redefine_test():
     expects_int(x) # $int
     x = str("Hello") # $str
     expects_string(x) # $str
+
+# ------------------------------------------------------------------------------
+# Tracking of self in methods
+# ------------------------------------------------------------------------------
+
+class Foo(object):
+
+    def meth1(self):
+        do_stuff(self)
+
+    def meth2(self): # $ MISSING: tracked_self
+        do_stuff(self) # $ MISSING: tracked_self
+
+    def meth3(self): # $ MISSING: tracked_self
+        do_stuff(self) # $ MISSING: tracked_self
+
+
+class Bar(Foo):
+
+    def meth1(self): # $ tracked_self
+        do_stuff(self) # $ tracked_self
+
+    def meth2(self):
+        do_stuff(self)
+
+    def meth3(self):
+        do_stuff(self)
+
+    def track_self(self): # $ tracked_self
+        self.meth1() # $ tracked_self
+        super().meth2()
+        super(Bar, self).foo3() # $ tracked_self

python/ql/test/experimental/dataflow/typetracking/tracked.ql

@@ -73,3 +73,32 @@ class TrackedStringTest extends InlineExpectationsTest {
     )
   }
 }
+
+DataFlow::Node tracked_self(TypeTracker t) {
+  t.start() and
+  exists(Function f |
+    f.isMethod() and
+    f.getName() = "track_self" and
+    result.(DataFlow::ParameterNode).getParameter() = f.getArg(0)
+  )
+  or
+  exists(TypeTracker t2 | result = tracked_self(t2).track(t2, t))
+}
+
+class TrackedSelfTest extends InlineExpectationsTest {
+  TrackedSelfTest() { this = "TrackedSelfTest" }
+
+  override string getARelevantTag() { result = "tracked_self" }
+
+  override predicate hasActualResult(Location location, string element, string tag, string value) {
+    exists(DataFlow::Node e, TypeTracker t |
+      e = tracked_self(t) and
+      // Module variables have no sensible location, and hence can't be annotated.
+      not e instanceof DataFlow::ModuleVariableNode and
+      tag = "tracked_self" and
+      location = e.getLocation() and
+      value = t.getAttr() and
+      element = e.toString()
+    )
+  }
+}