Add support for ends/startswith + tuple

Add support for startswith and endswith calls where
the argument is a tuple of prefixes/suffixes

Author: stsnel

README.md

@@ -40,7 +40,7 @@ in code, as well as overlapping execution traces.
 |Comparison | != (does not equal operator) | Full string | a != "foo" |
 |Comparison | in (part of operator) | String fragment | "foo" in a |
 |Comparison | in (match with collection) | List of full strings | a in ["foo","bar"] |
-|Call | string.startswith | String prefix | a.startswith("foo") |
-|Call | string.endswith | String suffix | a.endswith("foo") |
+|Call | string.startswith | String prefix, or tuple of string prefixes  | a.startswith("foo") |
+|Call | string.endswith | String suffix, or tuple of string suffixes  | a.endswith("foo") |
 |Call | string.find | String fragment | a.find("foo") |
 |Call | string.index | String fragment | a.index("foo") |

string_extractor/string_collector.py

@@ -28,12 +28,39 @@ def _isStringClass(self, c):
         else:
             return False
 
+    def _isTupleClass(self, c):
+        className = str(type(c))
+        return className in  ["<class '_ast.Tuple'>", "<class 'ast.Tuple'>"]
+
     def _getStringValue(self, c):
+        """Parses a single string value in the AST tree"""
         if str(type(c)) in ["<class '_ast.Str'>", "<class 'str'>"]:
             return c.s
         else:
             return c.value
 
+    def _getStringValuesAsList(self, c):
+        """Parses a single string value or tuple of strings in an AST
+           tree as a list of strings."""
+        def _get_single(c):
+            classname = str(type(c))
+            if classname in ["<class '_ast.Str'>", "<class 'str'>"]:
+                return [c.s]
+            elif classname in [ "<class '_ast.Constant'>", "<class 'ast.Constant'>" ]:
+                return [c.value]
+            else:
+                return []
+
+        classname = str(type(c))
+        if classname in ["<class '_ast.Str'>", "<class 'str'>",
+                         "<class '_ast.Constant'>", "<class 'ast.Constant'>"  ]:
+            return _get_single(c)
+        elif classname in ["<class '_ast.Tuple'>", "<class 'ast.Tuple'>"]:
+            result = []
+            for element in c.elts:
+                result.extend(_get_single(element))
+            return result
+
     def visit_Compare(self, node):
         opstype = str(type(node.ops[0]))
         if opstype in ["<class '_ast.Eq'>","<class '_ast.NotEq'>",
@@ -59,11 +86,11 @@ def visit_Compare(self, node):
                     if self._isStringClass(element):
                         self.fullStrings.add(self._getStringValue(element))
 
-    def visit_Call(self,node):
+    def visit_Call(self, node):
         try:
             attr = node.func.attr
+            args = node.args
             arg0 = node.args[0]
-            string = node.args[0].s
         except AttributeError:
             # Call without attribute or first argument
             # is not a regular string
@@ -72,10 +99,18 @@ def visit_Call(self,node):
             # Call without arguments
             return
 
-        if attr == "startswith" and self._isStringClass(arg0):
-            self.prefixes.add(self._getStringValue(arg0))
-        elif attr == "endswith" and self._isStringClass(arg0):
-            self.suffixes.add(self._getStringValue(arg0))
+        if attr == "startswith":
+            for arg in node.args:
+                if self._isStringClass(arg):
+                    self.prefixes.add(self._getStringValue(arg))
+                elif self._isTupleClass(arg):
+                    self.prefixes.update(self._getStringValuesAsList(arg))
+        elif attr == "endswith":
+            for arg in node.args:
+                if self._isStringClass(arg):
+                    self.suffixes.add(self._getStringValue(arg))
+                elif self._isTupleClass(arg):
+                    self.suffixes.update(self._getStringValuesAsList(arg))
         elif attr == "index"    and self._isStringClass(arg0):
             self.fragments.add(self._getStringValue(arg0))
         elif attr == "find"     and self._isStringClass(arg0):

tests/tests.py

@@ -118,18 +118,32 @@ def test_comparison_in(self):
         assert(output[0][0] == "FRAGMENT")
         assert(output[0][1] == "foo")
 
-    def test_call_startswith(self):
+    def test_call_startswith_single(self):
         output = self.extractor.getInterestingStrings('a.startswith("foo")')
         assert(len(output) == 1)
         assert(output[0][0] == "PREFIX")
         assert(output[0][1] == "foo")
 
-    def test_call_endswith(self):
+    def test_call_startswith_tuple(self):
+        output = self.extractor.getInterestingStrings('a.startswith( ("foo", "bar") )')
+        assert(len(output) == 2)
+        assert(output[0][0] == "PREFIX")
+        assert(output[1][0] == "PREFIX")
+        assert(sorted(list(map( lambda x : x[1], output ))) == ["bar", "foo"])
+
+    def test_call_endswith_single(self):
         output = self.extractor.getInterestingStrings('a.endswith("foo")')
         assert(len(output) == 1)
         assert(output[0][0] == "SUFFIX")
         assert(output[0][1] == "foo")
 
+    def test_call_endswith_tuple(self):
+        output = self.extractor.getInterestingStrings('a.endswith( ("foo", "bar") )')
+        assert(len(output) == 2)
+        assert(output[0][0] == "SUFFIX")
+        assert(output[1][0] == "SUFFIX")
+        assert(sorted(list(map( lambda x : x[1], output ))) == ["bar", "foo"])
+
     def test_call_index(self):
         output = self.extractor.getInterestingStrings('a.index("foo")')
         assert(len(output) == 1)