From 157258d8819707d42fa9c8d1c30b2ab129f9ae7c Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Tue, 10 Mar 2020 10:43:23 -0700
Subject: [PATCH] tags: Implement strip regex for docs processing

Co-Authored-By: Patrick Thomson <patrickt@users.noreply.github.com>
---
 cli/src/error.rs           |  6 ++++
 cli/src/tests/tags_test.rs | 16 +++++------
 tags/src/lib.rs            | 58 ++++++++++++++++++++++++++++++++------
 3 files changed, 64 insertions(+), 16 deletions(-)

diff --git a/cli/src/error.rs b/cli/src/error.rs
index 73dcb732..824bd92f 100644
--- a/cli/src/error.rs
+++ b/cli/src/error.rs
@@ -81,6 +81,12 @@ impl<'a> From<tree_sitter_highlight::Error> for Error {
     }
 }
 
+impl<'a> From<tree_sitter_tags::Error> for Error {
+    fn from(error: tree_sitter_tags::Error) -> Self {
+        Error::new(format!("{:?}", error))
+    }
+}
+
 impl From<serde_json::Error> for Error {
     fn from(error: serde_json::Error) -> Self {
         Error::new(error.to_string())
diff --git a/cli/src/tests/tags_test.rs b/cli/src/tests/tags_test.rs
index f8ad2527..2cc79749 100644
--- a/cli/src/tests/tags_test.rs
+++ b/cli/src/tests/tags_test.rs
@@ -10,13 +10,13 @@ fn test_tags_javascript() {
         ((function_definition
             name: (identifier) @name
             body: (block . (expression_statement (string) @doc))) @function
-         (set! replace @doc "(^['\s]*)|(['\s]*$)"))
+         (set! strip @doc "(^['\s]*)|(['\s]*$)"))
         (function_definition
             name: (identifier) @name) @function
         ((class_definition
             name: (identifier) @name
             body: (block . (expression_statement (string) @doc))) @class
-         (set! replace @doc "(^['\s]*)|(['\s]*$)"))
+         (set! strip @doc "(^['\s]*)|(['\s]*$)"))
         (class_definition
             name: (identifier) @name) @class
         (call
@@ -31,14 +31,14 @@ fn test_tags_javascript() {
         &tags_config,
         br#"
         class Customer:
-            """
+            '''
             Data about a customer
-            """
+            '''
 
             def age(self):
-                """
+                '''
                 Get the customer's age
-                """
+                '''
                 compute_age(self.id);
         }
         "#,
@@ -53,6 +53,6 @@ fn test_tags_javascript() {
         ]
     );
 
-    assert_eq!(tags[0].docs, Some("Data about a customer"));
-    assert_eq!(tags[1].docs, Some("Get the customer's age"));
+    assert_eq!(tags[0].docs.as_ref().unwrap(), "Data about a customer");
+    assert_eq!(tags[1].docs.as_ref().unwrap(), "Get the customer's age");
 }
diff --git a/tags/src/lib.rs b/tags/src/lib.rs
index b79a22df..577c4608 100644
--- a/tags/src/lib.rs
+++ b/tags/src/lib.rs
@@ -1,7 +1,8 @@
+use regex::Regex;
 use serde::{Serialize, Serializer};
 use std::collections::{hash_map, HashMap};
 use std::{ops, str};
-use tree_sitter::{Language, Node, Parser, Point, Query, QueryCursor, QueryError};
+use tree_sitter::{Language, Node, Parser, Query, QueryCursor, QueryError};
 
 /// Contains the data neeeded to compute tags for code written in a
 /// particular language.
@@ -15,6 +16,7 @@ pub struct TagsConfiguration {
     locals_pattern_index: usize,
     module_capture_index: Option<u32>,
     name_capture_index: Option<u32>,
+    doc_strip_regexes: Vec<Option<Regex>>,
 }
 
 pub struct TagsContext {
@@ -49,15 +51,17 @@ pub struct Tag<'a> {
     pub loc: Loc,
     pub name: &'a str,
     pub line: &'a str,
-    pub docs: Option<&'a str>,
+    pub docs: Option<String>,
+}
+
+#[derive(Debug)]
+pub enum Error {
+    Query(QueryError),
+    Regex(regex::Error),
 }
 
 impl TagsConfiguration {
-    pub fn new(
-        language: Language,
-        tags_query: &str,
-        locals_query: &str,
-    ) -> Result<Self, QueryError> {
+    pub fn new(language: Language, tags_query: &str, locals_query: &str) -> Result<Self, Error> {
         let query = Query::new(language, &format!("{}{}", tags_query, locals_query))?;
 
         let locals_query_offset = tags_query.len();
@@ -88,6 +92,23 @@ impl TagsConfiguration {
             *index = Some(i as u32);
         }
 
+        let doc_strip_regexes = (0..query.pattern_count())
+            .map(|pattern_index| {
+                let properties = query.property_settings(pattern_index);
+                for property in properties {
+                    if property.key.as_ref() == "strip"
+                        && property.capture_id.map(|id| id as u32) == doc_capture_index
+                    {
+                        if let Some(value) = &property.value {
+                            let regex = Regex::new(value.as_ref())?;
+                            return Ok(Some(regex));
+                        }
+                    }
+                }
+                return Ok(None);
+            })
+            .collect::<Result<Vec<_>, Error>>()?;
+
         Ok(TagsConfiguration {
             language,
             query,
@@ -98,6 +119,7 @@ impl TagsConfiguration {
             doc_capture_index,
             call_capture_index,
             name_capture_index,
+            doc_strip_regexes,
         })
     }
 }
@@ -158,7 +180,15 @@ impl TagsContext {
 
             let tag_from_node = |node: Node, kind: TagKind| -> Option<Tag> {
                 let name = str::from_utf8(&source[name_node?.byte_range()]).ok()?;
-                let docs = doc_node.and_then(|n| str::from_utf8(&source[n.byte_range()]).ok());
+                let docs = doc_node
+                    .and_then(|n| str::from_utf8(&source[n.byte_range()]).ok())
+                    .map(|s| {
+                        if let Some(regex) = &config.doc_strip_regexes[mat.pattern_index] {
+                            regex.replace_all(s, "").to_string()
+                        } else {
+                            s.to_string()
+                        }
+                    });
                 Some(Tag {
                     name,
                     line: "TODO",
@@ -233,3 +263,15 @@ impl From<tree_sitter::Point> for Pos {
         };
     }
 }
+
+impl From<regex::Error> for Error {
+    fn from(error: regex::Error) -> Self {
+        Error::Regex(error)
+    }
+}
+
+impl From<QueryError> for Error {
+    fn from(error: QueryError) -> Self {
+        Error::Query(error)
+    }
+}