lunr-search-index.js

#!/usr/bin/env node
import { fileURLToPath } from 'url'
import path from 'path'
import lunr from 'lunr'
import lunrStemmerSupport from 'lunr-languages/lunr.stemmer.support.js'
import tinyseg from 'lunr-languages/tinyseg.js'
import lunrJa from 'lunr-languages/lunr.ja.js'
import lunrEs from 'lunr-languages/lunr.es.js'
import lunrPt from 'lunr-languages/lunr.pt.js'
import fs from 'fs/promises'
import validateRecords from './validate-records.js'
import { compress } from '../../lib/search/compress.js'

const __dirname = path.dirname(fileURLToPath(import.meta.url))
lunrStemmerSupport(lunr)
tinyseg(lunr)
lunrJa(lunr)
lunrEs(lunr)
lunrPt(lunr)

export default class LunrIndex {
  constructor(name, records) {
    this.name = name

    // Add custom rankings
    this.records = records.map((record) => {
      return record
    })

    this.validate()

    return this
  }

  validate() {
    return validateRecords(this.name, this.records)
  }

  build() {
    const language = this.name.split('-').pop()
    const records = this.records

    this.index = lunr(function constructIndex() {
      // No arrow here!
      if (['ja', 'es', 'pt'].includes(language)) {
        this.use(lunr[language])
      }

      // By default Lunr considers the `-` character to be a word boundary.
      // This allows hyphens to be included in the search index.
      // If you change this, remember to make it match the indexing separator
      // in lib/search/lunr-search.js so the query is tokenized
      // identically to the way it was indexed.
      this.tokenizer.separator = /[\s]+/

      this.ref('objectID')
      this.field('url')
      this.field('breadcrumbs')
      this.field('headings', { boost: 3 })
      this.field('title', { boost: 5 })
      this.field('content')
      this.field('topics')

      this.metadataWhitelist = ['position']

      for (const record of records) {
        this.add(record)
      }
    })
  }

  toJSON() {
    this.build()
    return JSON.stringify(this.index, null, 2)
  }

  get recordsObject() {
    return Object.fromEntries(this.records.map((record) => [record.objectID, record]))
  }

  async write({
    outDirectory = path.posix.join(__dirname, '../../lib/search/indexes'),
    compressFiles = true,
  }) {
    this.build()

    // Write the parsed records
    await Promise.resolve(this.recordsObject)
      .then(JSON.stringify)
      .then((str) => (compressFiles ? compress(str) : str))
      .then((content) =>
        fs.writeFile(
          path.join(
            outDirectory,
            compressFiles ? `${this.name}-records.json.br` : `${this.name}-records.json`
          ),
          content
          // Do not set to 'utf8'
        )
      )

    // Write the index
    await Promise.resolve(this.index)
      .then(JSON.stringify)
      .then((str) => (compressFiles ? compress(str) : str))
      .then((content) =>
        fs.writeFile(
          path.join(outDirectory, compressFiles ? `${this.name}.json.br` : `${this.name}.json`),
          content
          // Do not set to 'utf8'
        )
      )
  }
}