From 8c78c91cc3f919d077246f8f14160e99b8fe28e3 Mon Sep 17 00:00:00 2001
From: Simon Cambier <simon.cambier@protonmail.com>
Date: Sat, 11 Feb 2023 15:43:01 +0100
Subject: [PATCH] #154 - filter by file extension

---
 src/globals.ts           |  1 +
 src/search/omnisearch.ts |  7 +++++++
 src/search/query.ts      | 21 +++++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/src/globals.ts b/src/globals.ts
index 29c7013..a9e77fc 100644
--- a/src/globals.ts
+++ b/src/globals.ts
@@ -8,6 +8,7 @@ export const regexLineSplit = /\r?\n|\r|((\.|\?|!)( |\r?\n|\r))/g
 export const regexYaml = /^---\s*\n(.*?)\n?^---\s?/ms
 export const regexStripQuotes = /^"|"$|^'|'$/g
 export const chsRegex = /[\u4e00-\u9fa5]/
+export const regexExtensions = /(?:^|\s)\.(\w+)/g
 
 export const excerptBefore = 100
 export const excerptAfter = 300
diff --git a/src/search/omnisearch.ts b/src/search/omnisearch.ts
index fad8c51..53b5746 100644
--- a/src/search/omnisearch.ts
+++ b/src/search/omnisearch.ts
@@ -175,6 +175,13 @@ export class Omnisearch {
       },
     })
 
+    // Filter query results to only keep files that match query.extensions (if any)
+    if (query.extensions.length) {
+      results = results.filter(r =>
+        query.extensions.some(e => r.id.endsWith(e))
+      )
+    }
+
     // If the query does not return any result,
     // retry but with a shorter prefix limit
     if (!results.length) {
diff --git a/src/search/query.ts b/src/search/query.ts
index bbe595c..c11701b 100644
--- a/src/search/query.ts
+++ b/src/search/query.ts
@@ -1,6 +1,7 @@
 import { settings } from '../settings'
 import { removeDiacritics, stripSurroundingQuotes } from '../tools/utils'
 import { parseQuery } from '../vendor/parse-query'
+import { regexExtensions } from '../globals'
 
 type QueryToken = {
   /**
@@ -20,8 +21,13 @@ type QueryToken = {
 export class Query {
   public segments: QueryToken[] = []
   public exclusions: QueryToken[] = []
+  public extensions: string[] = []
 
   constructor(text = '') {
+    // Extract & remove extensions from the query
+    this.extensions = this.extractExtensions(text)
+    text = this.removeExtensions(text)
+
     if (settings.ignoreDiacritics) text = removeDiacritics(text)
     const tokens = parseQuery(text.toLowerCase(), { tokenize: true })
     this.exclusions = tokens.exclude.text
@@ -59,4 +65,19 @@ export class Query {
       exact: stripped !== str,
     }
   }
+
+  /**
+   * Extracts an array of extensions like ".png" from a string
+   */
+  private extractExtensions(str: string): string[] {
+    const extensions = (str.match(regexExtensions) ?? []).map(o => o.trim())
+    if (extensions) {
+      return extensions.map(ext => ext.toLowerCase())
+    }
+    return []
+  }
+
+  private removeExtensions(str: string): string {
+    return str.replace(regexExtensions, '')
+  }
 }