__init__.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. import os
  2. import logging
  3. from mkdocs import utils
  4. from mkdocs.plugins import BasePlugin
  5. from mkdocs.config import config_options
  6. from mkdocs.contrib.search.search_index import SearchIndex
  7. log = logging.getLogger(__name__)
  8. base_path = os.path.dirname(os.path.abspath(__file__))
  9. class LangOption(config_options.OptionallyRequired):
  10. """ Validate Language(s) provided in config are known languages. """
  11. def lang_file_exists(self, lang):
  12. path = os.path.join(base_path, 'lunr-language', 'lunr.{}.js'.format(lang))
  13. return os.path.isfile(path)
  14. def run_validation(self, value):
  15. if isinstance(value, str):
  16. value = [value]
  17. elif not isinstance(value, (list, tuple)):
  18. raise config_options.ValidationError('Expected a list of language codes.')
  19. for lang in value:
  20. if lang != 'en' and not self.lang_file_exists(lang):
  21. raise config_options.ValidationError(
  22. '"{}" is not a supported language code.'.format(lang)
  23. )
  24. return value
  25. class SearchPlugin(BasePlugin):
  26. """ Add a search feature to MkDocs. """
  27. config_scheme = (
  28. ('lang', LangOption(default=['en'])),
  29. ('separator', config_options.Type(str, default=r'[\s\-]+')),
  30. ('min_search_length', config_options.Type(int, default=3)),
  31. ('prebuild_index', config_options.Choice((False, True, 'node', 'python'), default=False)),
  32. )
  33. def on_config(self, config, **kwargs):
  34. "Add plugin templates and scripts to config."
  35. if 'include_search_page' in config['theme'] and config['theme']['include_search_page']:
  36. config['theme'].static_templates.add('search.html')
  37. if not ('search_index_only' in config['theme'] and config['theme']['search_index_only']):
  38. path = os.path.join(base_path, 'templates')
  39. config['theme'].dirs.append(path)
  40. if 'search/main.js' not in config['extra_javascript']:
  41. config['extra_javascript'].append('search/main.js')
  42. return config
  43. def on_pre_build(self, config, **kwargs):
  44. "Create search index instance for later use."
  45. self.search_index = SearchIndex(**self.config)
  46. def on_page_context(self, context, **kwargs):
  47. "Add page to search index."
  48. self.search_index.add_entry_from_context(context['page'])
  49. def on_post_build(self, config, **kwargs):
  50. "Build search index."
  51. output_base_path = os.path.join(config['site_dir'], 'search')
  52. search_index = self.search_index.generate_search_index()
  53. json_output_path = os.path.join(output_base_path, 'search_index.json')
  54. utils.write_file(search_index.encode('utf-8'), json_output_path)
  55. if not ('search_index_only' in config['theme'] and config['theme']['search_index_only']):
  56. # Include language support files in output. Copy them directly
  57. # so that only the needed files are included.
  58. files = []
  59. if len(self.config['lang']) > 1 or 'en' not in self.config['lang']:
  60. files.append('lunr.stemmer.support.js')
  61. if len(self.config['lang']) > 1:
  62. files.append('lunr.multi.js')
  63. for lang in self.config['lang']:
  64. if (lang != 'en'):
  65. files.append('lunr.{}.js'.format(lang))
  66. for filename in files:
  67. from_path = os.path.join(base_path, 'lunr-language', filename)
  68. to_path = os.path.join(output_base_path, filename)
  69. utils.copy_file(from_path, to_path)