tantiny 0.2.2

data/lib/tantiny/query.rb

@@ -0,0 +1,165 @@
+# frozen_string_literal: true
+
+require "date"
+
+module Tantiny
+  class Query
+    TYPES = %i[
+      all empty term fuzzy_term
+      phrase regex range facet
+      smart prefix
+    ].freeze
+
+    DEFAULT_BOOST = 1.0
+    DEFAULT_FUZZY_DISTANCE = 1
+
+    class << self
+      def conjunction(*queries)
+        # @type var queries: Array[untyped]
+        queries.one? ? queries.first : __conjunction(queries)
+      end
+
+      def disjunction(*queries)
+        # @type var queries: Array[untyped]
+        queries.one? ? queries.first : __disjunction(queries)
+      end
+
+      def all_query(_index = nil)
+        __new_all_query
+      end
+
+      def empty_query(_index = nil)
+        __new_empty_query
+      end
+
+      def term_query(index, fields, term, **options)
+        allowed_fields = text_and_strings(index)
+        construct_query(index, :term, allowed_fields, fields, [term.to_s], **options)
+      end
+
+      def fuzzy_term_query(index, fields, term, distance = DEFAULT_FUZZY_DISTANCE, **options)
+        params = [term.to_s, distance.to_i]
+        allowed_fields = text_and_strings(index)
+        construct_query(index, :fuzzy_term, allowed_fields, fields, params, **options)
+      end
+
+      def phrase_query(index, fields, phrase, **options)
+        queries = [*fields].map do |f|
+          terms = index.schema.tokenizer_for(f).terms(phrase)
+          allowed_fields = index.schema.text_fields
+          construct_query(index, :phrase, allowed_fields, f, [terms], **options)
+        end
+
+        queries.empty? ? empty_query : disjunction(*queries)
+      end
+
+      def regex_query(index, fields, regex, **options)
+        allowed_fields = text_and_strings(index)
+        construct_query(index, :regex, allowed_fields, fields, [regex.to_s], **options)
+      end
+
+      def prefix_query(index, fields, prefix, **options)
+        regex_query(index, fields, Regexp.escape(prefix) + ".*", **options)
+      end
+
+      def range_query(index, fields, range, **options)
+        schema = index.schema
+
+        case range.first
+        when Integer
+          allowed_fields = schema.integer_fields
+          from, to = [range.min, range.max]
+        when Float
+          allowed_fields = schema.double_fields
+          from, to = [range.first, range.last]
+        when Date, DateTime
+          # @type var range: Range[Date | DateTime]
+          allowed_fields = schema.date_fields
+          from, to = [Helpers.timestamp(range.first), Helpers.timestamp(range.last)]
+        else
+          raise UnsupportedRange.new(range.first.class)
+        end
+
+        # @type var allowed_fields: Array[Symbol]
+        construct_query(index, :range, allowed_fields, fields, [from, to], **options)
+      end
+
+      def facet_query(index, field, path, **options)
+        allowed_fields = index.schema.facet_fields
+        construct_query(index, :facet, allowed_fields, field, [path], **options)
+      end
+
+      def smart_query(index, fields, query_string, **options)
+        fuzzy_distance = options[:fuzzy_distance]
+        boost_factor = options.fetch(:boost, DEFAULT_BOOST)
+
+        field_queries = [*fields].map do |field|
+          terms = index.schema.tokenizer_for(field).terms(query_string)
+
+          # See: https://github.com/soutaro/steep/issues/272
+          # @type block: nil | Query
+          next if terms.empty?
+
+          term_queries = terms.map do |term|
+            if fuzzy_distance.nil?
+              term_query(index, field, term)
+            else
+              fuzzy_term_query(index, field, term, fuzzy_distance)
+            end
+          end
+
+          # @type var terms: untyped
+          # @type var term_queries: untyped
+          last_term_query = prefix_query(index, field, terms.last) | term_queries.last
+
+          conjunction(last_term_query, *term_queries[0...-1])
+        end.compact
+
+        disjunction(*field_queries).boost(boost_factor)
+      end
+
+      private
+
+      # Can't use variadic argument `params` here due to:
+      # https://github.com/soutaro/steep/issues/480
+      def construct_query(index, query_type, allowed_fields, fields, params, **options)
+        queries = [*fields].map do |field|
+          supported = allowed_fields.include?(field)
+          raise UnsupportedField.new(field) unless supported
+
+          send("__new_#{query_type}_query", index, field.to_s, *params)
+        end
+
+        return empty_query if fields.empty?
+
+        disjunction(*queries).boost(options.fetch(:boost, DEFAULT_BOOST))
+      end
+
+      def text_and_strings(index)
+        index.schema.text_fields | index.schema.string_fields
+      end
+    end
+
+    def |(other)
+      raise ArgumentError.new("Not a #{self.class}.") unless other.is_a?(self.class)
+
+      self.class.disjunction(self, other)
+    end
+
+    def &(other)
+      raise ArgumentError.new("Not a #{self.class}.") unless other.is_a?(self.class)
+
+      self.class.conjunction(self, other)
+    end
+
+    def !
+      __negation
+    end
+
+    def boost(boost_factor)
+      return self if boost_factor == DEFAULT_BOOST
+
+      __boost(boost_factor.to_f)
+    end
+  end
+end

data/lib/tantiny/schema.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Tantiny
+  class Schema
+    attr_reader :default_tokenizer,
+      :id_field,
+      :text_fields,
+      :string_fields,
+      :integer_fields,
+      :double_fields,
+      :date_fields,
+      :facet_fields,
+      :field_tokenizers
+
+    def initialize(tokenizer, &block)
+      @default_tokenizer = tokenizer
+      @id_field = :id
+      @text_fields = []
+      @string_fields = []
+      @integer_fields = []
+      @double_fields = []
+      @date_fields = []
+      @facet_fields = []
+      @field_tokenizers = {}
+
+      instance_exec(&block)
+    end
+
+    def tokenizer_for(field)
+      field_tokenizers[field] || default_tokenizer
+    end
+
+    private
+
+    def id(key) = @id_field = key
+
+    def string(key) = @string_fields << key
+
+    def integer(key) = @integer_fields << key
+
+    def double(key) = @double_fields << key
+
+    def date(key) = @date_fields << key
+
+    def facet(key) = @facet_fields << key
+
+    def text(key, tokenizer: nil)
+      @field_tokenizers[key] = tokenizer if tokenizer
+
+      @text_fields << key
+    end
+  end
+end

data/lib/tantiny/tokenizer.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Tantiny
+  class Tokenizer
+    def self.default
+      new(:simple)
+    end
+
+    def self.new(kind, **options)
+      case kind
+      when :simple
+        __new_simple_tokenizer
+      when :stemmer
+        language = options[:language] || :en
+        __new_stemmer_tokenizer(language.to_s)
+      when :ngram
+        prefix_only = options.fetch(:prefix_only, false)
+        __new_ngram_tokenizer(options[:min], options[:max], prefix_only)
+      else
+        raise UnknownTokenizer.new(kind)
+      end
+    end
+
+    def terms(string)
+      __extract_terms(string)
+    end
+  end
+end

data/lib/tantiny/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Tantiny
+  VERSION = "0.2.2" # {x-release-please-version}
+end

data/lib/tantiny.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+require "ruby-next/language/setup"
+RubyNext::Language.setup_gem_load_path
+
+require "rutie"
+
+require "tantiny/version"
+require "tantiny/errors"
+require "tantiny/helpers"
+require "tantiny/schema"
+require "tantiny/tokenizer"
+require "tantiny/query"
+require "tantiny/index"
+
+require_relative "tantiny.so"
+
+module Tantiny
+end

data/sig/tantiny/errors.rbs

@@ -0,0 +1,20 @@
+module Tantiny
+  class TantivyError < StandardError
+  end
+
+  class UnknownField < StandardError
+    def initialize: () -> void
+  end
+
+  class UnknownTokenizer < StandardError
+    def initialize: (Symbol tokenizer_type) -> void
+  end
+
+  class UnsupportedRange < StandardError
+    def initialize: (Class range_type) -> void
+  end
+
+  class UnsupportedField < StandardError
+    def initialize: (Symbol field) -> void
+  end
+end

data/sig/tantiny/helpers.rbs

@@ -0,0 +1,6 @@
+
+module Tantiny
+  module Helpers
+    def self.timestamp: ((Date | DateTime) date) -> String
+  end
+end

data/sig/tantiny/index.rbs

@@ -0,0 +1,82 @@
+module Tantiny
+  class Index
+    DEFAULT_INDEX_SIZE: Integer
+    DEFAULT_LIMIT: Integer
+
+    def self.new: (
+      String path,
+      **untyped options
+    ) { (*untyped) -> void } -> Index
+
+    def self.__new: (
+      String path,
+      Integer index_size,
+      Tokenizer default_tokenizer,
+      Hash[String, Tokenizer] field_tokenizers,
+      Array[String] text_fields,
+      Array[String] string_fields,
+      Array[String] integer_fields,
+      Array[String] double_fields,
+      Array[String] date_fields,
+      Array[String] facet_fields
+    ) -> Index
+
+    attr_reader schema: Schema
+
+    def commit: () -> void
+    def reload: () -> void
+    def <<: (untyped document) -> void
+    def delete: (String id) -> void
+
+    def search: (
+      (Query | String) query,
+      ?limit: Integer,
+      **untyped smart_query_options
+    ) -> Array[String]
+
+    def all_query: () -> Query
+    def empty_query: () -> Query
+    def term_query: (fields fields, String term, **untyped options) -> Query
+    def fuzzy_term_query: (fields fields, String term, ?Integer distance, **untyped options) -> Query
+    def phrase_query: (fields fields, String phrase, **untyped options) -> Query
+    def regex_query: (fields fields, String regex, **untyped options) -> Query
+    def prefix_query: (fields fields, String prefix, **untyped options) -> Query
+    def facet_query: (Symbol field, String path, **untyped options) -> Query
+    def range_query: (fields fields, Range[numeric | date] range, **untyped options) -> Query
+    def smart_query: (fields fields, String query_string, **untyped options) -> Query
+
+    def __commit: () -> void
+    def __reload: () -> void
+
+    def __add_document: (
+      String id,
+      Hash[String, String] text_fields,
+      Hash[String, String] string_fields,
+      Hash[String, Integer] integer_fields,
+      Hash[String, Float] double_fields,
+      Hash[String, String] date_fields,
+      Hash[String, String] facet_fields
+    ) -> void
+
+    def __delete_document: (String id) -> void
+
+    def __search: (Query query, Integer limit) -> Array[String]
+
+    private
+
+    attr_writer schema: Schema
+
+    def slice_document: (
+      untyped document,
+      Array[Symbol] fields
+    ) { (untyped v) -> untyped } -> Hash[String, untyped]
+
+    def default_search: (
+      String query_string,
+      Integer limit,
+      ?fuzzy_distance: Integer
+    ) -> Array[String]
+
+    def resolve: (untyped document, Symbol field) -> untyped
+  end
+end

data/sig/tantiny/query.rbs

@@ -0,0 +1,135 @@
+module Tantiny
+  class Query
+    TYPES: Array[Symbol]
+    DEFAULT_BOOST: Float
+    DEFAULT_FUZZY_DISTANCE: Integer
+
+    def self.disjunction: (*Query queries) -> Query
+    def self.conjunction: (*Query queries) -> Query
+
+    def self.all_query: (?Index _index) -> Query
+    def self.empty_query: (?Index _index) -> Query
+
+    def self.term_query: (
+      Index index,
+      fields fields,
+      String term,
+      **untyped options
+    ) -> Query
+
+    def self.fuzzy_term_query: (
+      Index index,
+      fields fields,
+      String term,
+      ?Integer distance,
+      **untyped options
+    ) -> Query
+
+    def self.phrase_query: (
+      Index index,
+      fields fields,
+      String phrase,
+      **untyped options
+    ) -> Query
+
+    def self.regex_query: (
+      Index index,
+      fields fields,
+      String regex,
+      **untyped options
+    ) -> Query
+
+    def self.prefix_query: (
+      Index index,
+      fields fields,
+      String prefix,
+      **untyped options
+    ) -> Query
+
+    def self.facet_query: (
+      Index index,
+      Symbol field,
+      String path,
+      **untyped options
+    ) -> Query
+
+    def self.range_query: (
+      Index index,
+      fields fields,
+      Range[numeric | date] range,
+      **untyped options
+    ) -> Query
+
+    def self.smart_query: (
+      Index index,
+      fields fields,
+      String query_string,
+      **untyped options
+    ) -> Query
+
+    def self.__new_all_query: () -> Query
+    def self.__new_empty_query: () -> Query
+
+    def self.__new_term_query: (
+      Index index,
+      String field,
+      String term
+    ) -> Query
+
+    def self.__new_fuzzy_term_query: (
+      Index index,
+      String field,
+      String term,
+      Integer distance
+    ) -> Query
+
+    def self.__new_regex_query: (
+      Index index,
+      String field,
+      String regex
+    ) -> Query
+
+    def self.__new_range_query: (
+      Index index,
+      String field,
+      untyped from,
+      untyped to
+    ) -> Query
+
+    def self.__new_phrase_query: (
+      Index index,
+      String field,
+      Array[String] terms
+    ) -> Query
+
+    def self.__new_facet_query: (
+      Index index,
+      String field,
+      String path
+    ) -> Query
+
+    def self.__disjunction: (Array[Query] queries) -> Query
+    def self.__conjunction: (Array[Query] queries) -> Query
+
+    def |: (Query query) -> Query
+    def &: (Query query) -> Query
+    def !: () -> Query
+    def boost: (numeric boost_factor) -> Query
+
+    def __negation: () -> Query
+    def __boost: (Float boost_factor) -> Query
+
+    private
+    
+    def self.construct_query: (
+      Index index,
+      Symbol query_type,
+      Array[Symbol] allowed_fields,
+      fields fields,
+      Array[untyped] params,
+      **untyped options
+    ) -> Query
+
+    def self.text_and_strings: (Index index) -> Array[Symbol]
+  end
+end

data/sig/tantiny/schema.rbs

@@ -0,0 +1,26 @@
+module Tantiny
+  class Schema
+    attr_reader default_tokenizer: Tokenizer
+    attr_reader id_field: Symbol
+    attr_reader text_fields: Array[Symbol]
+    attr_reader string_fields: Array[Symbol]
+    attr_reader integer_fields: Array[Symbol]
+    attr_reader double_fields: Array[Symbol]
+    attr_reader date_fields: Array[Symbol]
+    attr_reader facet_fields: Array[Symbol]
+    attr_reader field_tokenizers: Hash[Symbol, Tokenizer]
+
+    def initialize: (Tokenizer tokenizer) { (*untyped) -> void } -> void
+    def tokenizer_for: (Symbol field) -> Tokenizer
+
+    private
+
+    def id: (Symbol key) -> void
+    def text: (Symbol key, ?tokenizer: Tokenizer) -> void
+    def string: (Symbol key) -> void
+    def integer: (Symbol key) -> void
+    def double: (Symbol key) -> void
+    def date: (Symbol key) -> void
+    def facet: (Symbol key) -> void
+  end
+end

data/sig/tantiny/tokenizer.rbs

@@ -0,0 +1,25 @@
+module Tantiny
+  class Tokenizer
+    def self.default: () -> Tokenizer
+
+    def self.new: (Symbol kind, **untyped options) -> Tokenizer
+
+    def self.__new_ngram_tokenizer: (
+      Integer min,
+      Integer max,
+      bool prefix_only
+    ) -> Tokenizer
+
+    def self.__new_stemmer_tokenizer: (
+      String locale_code
+    ) -> Tokenizer
+
+    def self.__new_simple_tokenizer: () -> Tokenizer
+
+    public
+
+    def terms: (String string) -> Array[String]
+
+    def __extract_terms: (String string) -> Array[String]
+  end
+end

data/sig/tantiny/version.rbs

@@ -0,0 +1,3 @@
+module Tantiny
+  VERSION: String
+end

data/sig/tantiny.rbs

@@ -0,0 +1,5 @@
+module Tantiny
+    type date = Date | DateTime
+    type numeric = Integer | Float
+    type fields = Array[Symbol] | Symbol
+end