Python: Only include relevant YAML in printAst.ql

Author: hvitved

python/ql/lib/printAst.ql

@@ -25,4 +25,9 @@ class PrintAstConfigurationOverride extends PrintAstConfiguration {
     super.shouldPrint(e, l) and
     l.getFile() = getFileBySourceArchiveName(selectedSourceFile())
   }
+
+  override predicate shouldPrintYaml(YamlNode y, Location l) {
+    super.shouldPrintYaml(y, l) and
+    l.getFile() = getFileBySourceArchiveName(selectedSourceFile())
+  }
 }

python/ql/lib/semmle/python/PrintAst.qll

@@ -26,12 +26,22 @@ class PrintAstConfiguration extends TPrintAstConfiguration {
    * By default it checks whether the `AstNode` `e` belongs to `Location` `l`.
    */
   predicate shouldPrint(AstNode e, Location l) { l = e.getLocation() }
+
+  /**
+   * Controls whether the `YamlNode` should be considered for AST printing.
+   * By default it checks whether the `YamlNode` `y` belongs to `Location` `l`.
+   */
+  predicate shouldPrintYaml(YamlNode y, Location l) { l = y.getLocation() }
 }
 
 private predicate shouldPrint(AstNode e, Location l) {
   exists(PrintAstConfiguration config | config.shouldPrint(e, l))
 }
 
+private predicate shouldPrintYaml(YamlNode y, Location l) {
+  exists(PrintAstConfiguration config | config.shouldPrintYaml(y, l))
+}
+
 /** Holds if the given element does not need to be rendered in the AST. */
 private predicate isNotNeeded(AstNode el) {
   el.isArtificial()
@@ -55,8 +65,11 @@ private newtype TPrintAstNode =
     not list = any(Module mod).getBody() and
     not forall(AstNode child | child = list.getAnItem() | isNotNeeded(child))
   } or
-  TYamlNode(YamlNode node) or
-  TYamlMappingNode(YamlMapping mapping, int i) { exists(mapping.getKeyNode(i)) }
+  TYamlNode(YamlNode node) { shouldPrintYaml(node, _) } or
+  TYamlMappingNode(YamlMapping mapping, int i) {
+    shouldPrintYaml(mapping, _) and
+    exists(mapping.getKeyNode(i))
+  }
 
 /**
  * A node in the output tree.