Python: Add tests with redefinition of fields/variables

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

@@ -1,15 +1,31 @@
+class SomeClass:
+    pass
+
 def simple_read_write():
-    x = object() # $tracked=foo
+    x = SomeClass() # $tracked=foo
     x.foo = tracked # $tracked $tracked=foo
     y = x.foo # $tracked=foo $tracked
     do_stuff(y) # $tracked
 
 def foo():
-    x = object() # $tracked=attr
+    x = SomeClass() # $tracked=attr
     bar(x) # $tracked=attr
     x.attr = tracked # $tracked=attr $tracked
     baz(x) # $tracked=attr
 
 def bar(x): # $tracked=attr
     z = x.attr # $tracked $tracked=attr
-    do_stuff(z) # $tracked
+    do_stuff(z) # $tracked
+
+def expects_int(x): # $int=field $f+:str=field
+    do_int_stuff(x.field) # $int $f+:str $int=field $f+:str=field
+
+def expects_string(x): # $f+:int=field $str=field
+    do_string_stuff(x.field) # $f+:int $str $f+:int=field $str=field
+
+def test_incompatible_types():
+    x = SomeClass() # $int,str=field
+    x.field = int(5) # $int=field $f+:str=field $int $f+:str
+    expects_int(x) # $int=field $f+:str=field
+    x.field = str("Hello") # $f+:int=field $str=field $f+:int $str
+    expects_string(x) # $f+:int=field $str=field

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

@@ -47,3 +47,15 @@ def global_var_write_test():
     x = tracked # $tracked
     write_g(x) # $tracked
     use_g()
+
+def expects_int(x): # $int
+    do_int_stuff(x) # $int
+
+def expects_string(x): # $str
+    do_string_stuff(x) # $str
+
+def redefine_test():
+    x = int(5) # $int
+    expects_int(x) # $int
+    x = str("Hello") # $str
+    expects_string(x) # $str

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

@@ -19,7 +19,53 @@ class TrackedTest extends InlineExpectationsTest {
       e = tracked(t) and
       tag = "tracked" and
       location = e.getLocation() and
-      value = t.getProp() and
+      value = t.getAttr() and
+      element = e.toString()
+    )
+  }
+}
+
+Node int_type(TypeTracker t) {
+  t.start() and
+  result.asCfgNode() = any(CallNode c | c.getFunction().(NameNode).getId() = "int")
+  or
+  exists(TypeTracker t2 | result = int_type(t2).track(t2, t))
+}
+
+Node string_type(TypeTracker t) {
+  t.start() and
+  result.asCfgNode() = any(CallNode c | c.getFunction().(NameNode).getId() = "str")
+  or
+  exists(TypeTracker t2 | result = string_type(t2).track(t2, t))
+}
+
+class TrackedIntTest extends InlineExpectationsTest {
+  TrackedIntTest() { this = "TrackedIntTest" }
+
+  override string getARelevantTag() { result = "int" }
+
+  override predicate hasActualResult(Location location, string element, string tag, string value) {
+    exists(Node e, TypeTracker t |
+      e = int_type(t) and
+      tag = "int" and
+      location = e.getLocation() and
+      value = t.getAttr() and
+      element = e.toString()
+    )
+  }
+}
+
+class TrackedStringTest extends InlineExpectationsTest {
+  TrackedStringTest() { this = "TrackedStringTest" }
+
+  override string getARelevantTag() { result = "str" }
+
+  override predicate hasActualResult(Location location, string element, string tag, string value) {
+    exists(Node e, TypeTracker t |
+      e = string_type(t) and
+      tag = "str" and
+      location = e.getLocation() and
+      value = t.getAttr() and
       element = e.toString()
     )
   }