<commit_before>require_relative 'antidote/type_checker'
require_relative "antidote/variable_type_error"
require_relative "antidote/return_type_error"

module Antidote
  module ClassMethods
    @@type_constraints = Hash.new { |h, k| h[k] = [] }

    def annotate(type_constraints, method_name)
      add_type_constraints(method_name, type_constraints)
      define_type_check_for_instance_method(method_name)
    end

    def annotate_class_method(type_constraints, method_name)
      add_type_constraints("self.#{method_name}", type_constraints)
      define_type_check_for_class_method("self.#{method_name}")
    end

    private

    def add_type_constraints(method_name, type_constraints)
      type_constraints.first.each do |variable|
        @@type_constraints[method_name] << variable
      end

      @@type_constraints[method_name] << {__returns: type_constraints.last}
    end

    def define_type_check_for_class_method(method_name)
      method_name = method_name.split(".").last
      new_name_for_old_function = "#{method_name}_old".to_sym
      self.singleton_class.send(:alias_method, new_name_for_old_function,
                                method_name)

      define_singleton_method(method_name) do |*args|
        args.each_with_index do |argument_value, index|
          actual_type   = argument_value.class
          expected_type =
            @@type_constraints["self.#{method_name}"].first.values[index]
          variable_name =
            @@type_constraints["self.#{method_name}"].first.keys[index]
          @@return_type = @@type_constraints["self.#{method_name}"].last.values.first

          unless actual_type == expected_type
            raise Antidote::VariableTypeError,
              [actual_type, expected_type, variable_name, "self.#{method_name}"]
          end
        end

        if (klass = send(new_name_for_old_function, *args).class) != @@return_type
          raise Antidote::ReturnTypeError, [klass, @@return_type, "self.#{method_name}"]
        else
          self.send(new_name_for_old_function, *args)
        end
      end
    end

    def define_type_check_for_instance_method(method_name)
      new_name_for_old_function = "#{method_name}_old".to_sym
      alias_method(new_name_for_old_function, method_name)

      define_method(method_name) do |*args|
        args.each_with_index do |argument_value, index|
          actual_type   = argument_value.class
          expected_type = @@type_constraints[method_name].first.values[index]
          variable_name = @@type_constraints[method_name].first.keys[index]

          @@return_type = @@type_constraints[method_name].last.values.first

          unless actual_type == expected_type
            raise Antidote::VariableTypeError,
              [actual_type, expected_type, variable_name, method_name]
          end
        end

        if (klass = send(new_name_for_old_function, *args).class) != @@return_type
          raise Antidote::ReturnTypeError, [klass, @@return_type, method_name]
        else
          send(new_name_for_old_function, *args)
        end
      end
    end
  end

  class << self
    def included(klass)
      klass.extend(ClassMethods)
    end
  end
end
<commit_msg>Update syntax for annotations.<commit_after>require_relative "antidote/variable_type_error"
require_relative "antidote/return_type_error"

module Antidote
  class << self
    def included(klass)
      klass.extend(ClassMethods)
    end
  end

  module ClassMethods
    @@type_constraints = Hash.new { |h, k| h[k] = [] } 

    def annotate(*variables, return_type, &method_name)
      method_name = method_name.call.to_s
      add_type_constraints(method_name, variables, return_type)
      redefine_instance_method(method_name)
    end

    def annotate_class_method(*variables, return_type, &method_name)
      method_name = "self.#{method_name.call.to_s}"
      add_type_constraints(method_name, variables, return_type)
      redefine_class_method(method_name)
    end

    private

    def add_type_constraints(method_name, variables, return_type)
      variables.each { |variable| @@type_constraints[method_name] << variable }
      @@type_constraints[method_name] << {__returns: return_type}
    end

    def redefine_class_method(method_name)
      method_name = method_name.split(".").last
      new_name_for_old_method = "#{method_name}_old".to_sym
      self.singleton_class.send(:alias_method, new_name_for_old_method, method_name)

      define_singleton_method(method_name) do |*args|
        args.each_with_index do |argument, index|
          class_method = "self.#{method_name}"
          actual = argument.class
          expected = @@type_constraints[class_method][index]
          parameters = self.method(new_name_for_old_method).parameters.map(&:last)
          variable_name = parameters[index]
          @__return_type = @@type_constraints[class_method].last[:__returns]

          unless actual == expected
            raise VariableTypeError, [actual, expected, variable_name, class_method]
          end
        end

        if (klass = self.send(new_name_for_old_method, *args).class) != @__return_type 
          raise ReturnTypeError, [klass, @__return_type, class_method]
        else
          self.send(new_name_for_old_method, *args)
        end
      end
    end

    def redefine_instance_method(method_name)
      new_name_for_old_method = "#{method_name}_old".to_sym
      alias_method(new_name_for_old_method, method_name)

      define_method(method_name) do |*args|
        args.each_with_index do |argument, index|
          actual = argument.class
          expected = @@type_constraints[method_name][index]
          parameters = self.method(new_name_for_old_method).parameters.map(&:last)
          variable_name = parameters[index]
          @__return_type = @@type_constraints[method_name].last[:__returns]

          unless actual == expected
            raise VariableTypeError, [actual, expected, variable_name, method_name]
          end
        end

        if (klass = self.send(new_name_for_old_method, *args).class) != @__return_type 
          raise ReturnTypeError, [klass, @__return_type, method_name]
        else
          self.send(new_name_for_old_method, *args)
        end
      end
    end
  end
end

{'fraction_non_alphanumeric': 0.09209896249002394, 'fraction_numerical': 0.0, 'mean_word_length': 3.2366463826910072, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '58019', 'n_tokens_mistral': 1958, 'n_tokens_neox': 1918, 'n_words': 316}