Python: improve location computation

python/ql/lib/semmle/python/AstExtended.qll

@@ -154,6 +154,24 @@ class StringPart extends StringPart_, AstNode {
   override string toString() { result = StringPart_.super.toString() }
 
   override Location getLocation() { result = StringPart_.super.getLocation() }
+
+  /** Holds if the content of string `StringPart` is surrounded by `prefix` and `quote`. */
+  predicate context(string prefix, string quote) {
+    exists(int occurrenceOffset |
+      quote = this.getText().regexpFind("\"{3}|\"{1}|'{3}|'{1}", 0, occurrenceOffset) and
+      prefix = this.getText().prefix(occurrenceOffset + quote.length())
+    )
+  }
+
+  /**
+   * Gets the length of the content, that is the text between the prefix and the quote.
+   * See `context` for obtaining the prefix and the quote.
+   */
+  int getContentLenght() {
+    exists(string prefix, string quote | this.context(prefix, quote) |
+      result = this.getText().length() - prefix.length() - quote.length()
+    )
+  }
 }
 
 class StringPartList extends StringPartList_ { }

python/ql/lib/semmle/python/regexp/RegexTreeView.qll

@@ -223,16 +223,53 @@ module Impl implements RegexTreeViewSig {
      */
     Location getLocation() { result = re.getLocation() }
 
+    /** Gets the accumulated length of string parts with lower index than `index`, if any. */
+    private int getPartOffset(int index) {
+      index = 0 and result = 0
+      or
+      index > 0 and
+      exists(int previousOffset | previousOffset = this.getPartOffset(index - 1) |
+        result =
+          previousOffset + re.(StrConst).getImplicitlyConcatenatedPart(index - 1).getContentLenght()
+      )
+    }
+
+    /**
+     * Gets the `StringPart` in which this `RegExpTerm` resides, if any.
+     * `localOffset` will be the offset of this `RegExpTerm` inside `result`.
+     */
+    StringPart getPart(int localOffset) {
+      exists(int index, int prefixLength | index = max(int i | this.getPartOffset(i) < start) |
+        result = re.(StrConst).getImplicitlyConcatenatedPart(index) and
+        exists(string prefix | result.context(prefix, _) | prefixLength = prefix.length()) and
+        // Example:
+        // re.compile('...' r"""...this..""")
+        // - `start` is the offset from `(` to `this` as counted after concatenating all parts.
+        // - we subtract the lenght of the previous `StringPart`s, `'...'`, to know how far into this `StringPart` we go.
+        // - as the prefix 'r"""' is part of the `StringPart`, `this` is found that much further in.
+        localOffset = start - this.getPartOffset(index) + prefixLength
+      )
+    }
+
     /** Holds if this term is found at the specified location offsets. */
     predicate hasLocationInfo(
       string filepath, int startline, int startcolumn, int endline, int endcolumn
     ) {
+      not exists(this.getPart(_)) and
       exists(int re_start, int prefix_len | prefix_len = re.getPrefix().length() |
         re.getLocation().hasLocationInfo(filepath, startline, re_start, endline, _) and
         startcolumn = re_start + start + prefix_len and
         endcolumn = re_start + end + prefix_len - 1
         /* inclusive vs exclusive */
       )
+      or
+      exists(StringPart part, int localOffset | part = this.getPart(localOffset) |
+        filepath = part.getLocation().getFile().getAbsolutePath() and
+        startline = part.getLocation().getStartLine() and
+        startcolumn = part.getLocation().getStartColumn() + localOffset and
+        endline = startline and
+        endcolumn = (end - start) + startcolumn
+      )
     }
 
     /** Gets the file in which this term is found. */

python/ql/test/library-tests/regexparser/locations.py

@@ -50,25 +50,25 @@ br'''[this] is a test'''
 )
 
 # plain string with multiple parts
-re.compile( # $ location=1:2 SPURIOUS:location=1:23 MISSING:location=1:26
+re.compile( # $ location=1:2 location=1:26
 '[this] is a test' ' and [this] is another test'
 )
 
 # plain string with multiple parts across lines
-re.compile( # $ location=1:2 SPURIOUS:location=1:23 MISSING:location=2:7
+re.compile( # $ location=1:2 location=2:7
 '[this] is a test'
 ' and [this] is another test'
 )
 
 # plain string with multiple parts across lines and comments
-re.compile( # $ location=1:2 SPURIOUS:location=1:23 MISSING:location=3:7
+re.compile( # $ location=1:2 location=3:7
 '[this] is a test'
 # comment
 ' and [this] is another test'
 )
 
 # multiple parts of different kinds
-re.compile( # $ location=1:2 SPURIOUS:location=1:23 location=1:50 location=1:81 MISSING:location=1:28 location=2:11 location=3:8
+re.compile( # $ location=1:2 location=1:28 location=2:11 location=3:8
 '[this] is a test' ''' and [this] is another test'''
 br""" and [this] is yet another test"""
 r' and [this] is one more'