rbs_goose 0.1.0

data/lib/rbs_goose/configuration.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'openai'
+require 'langchain'
+
+module RbsGoose
+  class Configuration
+    def initialize(&block)
+      self.instruction = default_instruction
+      self.example_groups = default_example_groups
+      self.template_class = default_template_class
+      instance_eval(&block) if block_given?
+    end
+
+    attr_accessor :llm, :instruction, :example_groups, :template_class
+
+    def use_open_ai(open_ai_access_token, default_options: {})
+      @llm = ::Langchain::LLM::OpenAI.new(
+        api_key: open_ai_access_token,
+        default_options: {
+          completion_model_name: 'gpt-3.5-turbo-0613',
+          chat_completion_model_name: 'gpt-3.5-turbo-0613'
+        }.merge(default_options)
+      )
+    end
+
+    def template
+      @template ||= template_class.new(instruction: instruction, example_groups: example_groups)
+    end
+
+    private
+
+    def default_template_class
+      Templates::DefaultTemplate
+    end
+
+    def default_instruction
+      <<~INSTRUCTION
+        Act as Ruby type inferrer.
+        When ruby source codes and RBS type signatures are given, refine each RBS type signatures. Each file should be split in markdown code format.
+        Use class names, variable names, etc., to infer type.
+      INSTRUCTION
+    end
+
+    def default_example_groups
+      [RbsGoose::IO::ExampleGroup.default_examples[:rbs_samples]]
+    end
+  end
+end

data/lib/rbs_goose/examples/rbs_samples/Steepfile

@@ -0,0 +1,6 @@
+target :lib do
+  # signature "sig"
+  signature "refined/sig"
+
+  check "lib"
+end

data/lib/rbs_goose/examples/rbs_samples/lib/email.rb

@@ -0,0 +1,16 @@
+class Email
+  # @dynamic address
+  attr_reader :address
+
+  def initialize(address:)
+    @address = address
+  end
+
+  def ==(other)
+    other.is_a?(self.class) && other.address == address
+  end
+
+  def hash
+    self.class.hash ^ address.hash
+  end
+end

data/lib/rbs_goose/examples/rbs_samples/lib/person.rb

@@ -0,0 +1,19 @@
+class Person
+  # @dynamic name, contacts
+  attr_reader :name
+  attr_reader :contacts
+
+  def initialize(name:)
+    @name = name
+    @contacts = []
+  end
+
+  def guess_country()
+    contacts.map do |contact|
+      case contact
+      when Phone
+        contact.country
+      end
+    end.compact.first
+  end
+end

data/lib/rbs_goose/examples/rbs_samples/lib/phone.rb

@@ -0,0 +1,22 @@
+class Phone
+  # @dynamic country, number
+  attr_reader :country, :number
+
+  def initialize(country:, number:)
+    @country = country
+    @number = number
+  end
+
+  def ==(other)
+    if other.is_a?(Phone)
+      # @type var other: Phone
+      other.country == country && other.number == number
+    else
+      false
+    end
+  end
+
+  def hash
+    self.class.hash ^ country.hash ^ number.hash
+  end
+end

data/lib/rbs_goose/examples/rbs_samples/refined/sig/email.rbs

@@ -0,0 +1,11 @@
+class Email
+  @address: String
+
+  attr_reader address: String
+
+  def initialize: (address: String) -> void
+
+  def ==: (untyped other) -> bool
+
+  def hash: () -> Integer
+end

data/lib/rbs_goose/examples/rbs_samples/refined/sig/person.rbs

@@ -0,0 +1,13 @@
+class Person
+  @name: String
+
+  @contacts: Array[Email | Phone]
+
+  attr_reader name: String
+
+  attr_reader contacts: Array[Email | Phone]
+
+  def initialize: (name: String) -> void
+
+  def guess_country: () -> (String | nil)
+end

data/lib/rbs_goose/examples/rbs_samples/refined/sig/phone.rbs

@@ -0,0 +1,15 @@
+class Phone
+  @country: String
+
+  @number: String
+
+  attr_reader country: String
+
+  attr_reader number: String
+
+  def initialize: (country: String, number: String) -> void
+
+  def ==: (untyped other) -> (bool | nil)
+
+  def hash: () -> Integer
+end

data/lib/rbs_goose/examples/rbs_samples/sig/email.rbs

@@ -0,0 +1,11 @@
+class Email
+  @address: untyped
+
+  attr_reader address: untyped
+
+  def initialize: (address: untyped) -> void
+
+  def ==: (untyped other) -> untyped
+
+  def hash: () -> untyped
+end

data/lib/rbs_goose/examples/rbs_samples/sig/person.rbs

@@ -0,0 +1,13 @@
+class Person
+  @name: untyped
+
+  @contacts: untyped
+
+  attr_reader name: untyped
+
+  attr_reader contacts: untyped
+
+  def initialize: (name: untyped) -> void
+
+  def guess_country: () -> untyped
+end

data/lib/rbs_goose/examples/rbs_samples/sig/phone.rbs

@@ -0,0 +1,15 @@
+class Phone
+  @country: untyped
+
+  @number: untyped
+
+  attr_reader country: untyped
+
+  attr_reader number: untyped
+
+  def initialize: (country: untyped, number: untyped) -> void
+
+  def ==: (untyped other) -> (untyped | nil)
+
+  def hash: () -> untyped
+end

data/lib/rbs_goose/io/example.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module RbsGoose
+  module IO
+    class Example
+      class << self
+        def from_path(ruby_path:, rbs_path:, refined_rbs_dir:, base_path:)
+          Example.new(
+            typed_ruby: TypedRuby.from_path(ruby_path: ruby_path, rbs_path: rbs_path, base_path: base_path),
+            refined_rbs: File.new(path: rbs_path, base_path: ::File.join(base_path, refined_rbs_dir))
+          )
+        end
+      end
+
+      def initialize(typed_ruby:, refined_rbs:)
+        @typed_ruby = typed_ruby
+        @refined_rbs = refined_rbs
+      end
+
+      def to_h
+        { typed_ruby: typed_ruby, refined_rbs: refined_rbs }
+      end
+
+      attr_reader :typed_ruby, :refined_rbs
+    end
+  end
+end

data/lib/rbs_goose/io/example_group.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module RbsGoose
+  module IO
+    class ExampleGroup < Array
+      class << self
+        def load_from(base_path, code_dir: 'lib', sig_dir: 'sig', refined_dir: 'refined')
+          new.tap do |group|
+            Dir.glob('**/*.rb', base: ::File.join(base_path, code_dir)).each do |path|
+              group << Example.from_path(
+                ruby_path: ::File.join(code_dir, path),
+                rbs_path: to_rbs_path(path, sig_dir),
+                refined_rbs_dir: refined_dir,
+                base_path: base_path
+              )
+            end
+          end
+        end
+
+        def default_examples
+          example_dir = ::File.join(__dir__.to_s, '../examples')
+          @default_examples ||= Dir.glob('*', base: example_dir).to_h do |dir|
+            [dir.to_sym, load_from(::File.join(example_dir, dir))]
+          end
+        end
+
+        private
+
+        def to_rbs_path(path, sig_dir)
+          ::File.join(sig_dir, "#{path}s")
+        end
+      end
+
+      def to_target_group
+        TargetGroup.new.tap do |g|
+          each { g << _1.typed_ruby }
+        end
+      end
+
+      def to_refined_rbs_list
+        map(&:refined_rbs)
+      end
+    end
+  end
+end

data/lib/rbs_goose/io/file.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module RbsGoose
+  module IO
+    class File
+      MARKDOWN_REGEXP = /\A```(?<type>ruby|rbs):(?<path>.+)\n(?<content>[\s\S]+)```\Z/
+
+      class << self
+        def from_markdown(markdown)
+          parsed = markdown.match(MARKDOWN_REGEXP)
+          raise ArgumentError, "Ruby or RBS Markdown parsing failed.\n#{markdown}" unless parsed
+
+          new(path: parsed[:path].to_s, content: parsed[:content].to_s)
+        end
+      end
+
+      def initialize(path:, content: nil, base_path: nil)
+        @path = path
+        @base_path = base_path
+        @content = content&.strip
+        load_content if @content.nil?
+      end
+
+      def load_content
+        @content = ::File.read(absolute_path).strip
+      end
+
+      def absolute_path
+        base_path ? ::File.join(base_path, path) : path
+      end
+
+      def type
+        @type ||= case path
+                  in /\.rb\z/
+                    :ruby
+                  in /\.rbs\z/
+                    :rbs
+                  else
+                    raise ArgumentError, "Unknown file type: #{path}"
+                  end
+      end
+
+      def to_s
+        "```#{type}:#{path}\n#{content}\n```\n"
+      end
+
+      def content=(content)
+        @content = content.strip
+      end
+
+      def write
+        ::File.write(path, content)
+      end
+
+      def ==(other)
+        self.class == other.class &&
+          @path == other.path &&
+          @content == other.content
+      end
+
+      attr_reader :path, :content, :base_path
+    end
+  end
+end

data/lib/rbs_goose/io/target_group.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RbsGoose
+  module IO
+    class TargetGroup < Array
+      class << self
+        def load_from(base_path, code_dir: 'lib', sig_dir: 'sig')
+          new.tap do |group|
+            Dir.glob('**/*.rb', base: ::File.join(base_path, code_dir)).each do |path|
+              group << TypedRuby.from_path(
+                ruby_path: ::File.join(code_dir, path),
+                rbs_path: ::File.join(sig_dir, "#{path}s"),
+                base_path: base_path
+              )
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rbs_goose/io/typed_ruby.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module RbsGoose
+  module IO
+    class TypedRuby
+      class << self
+        def from_path(ruby_path:, rbs_path:, base_path:)
+          ruby = File.new(path: ruby_path, base_path: base_path)
+          rbs = begin
+            File.new(path: rbs_path, base_path: base_path)
+          rescue StandardError
+            nil
+          end
+          new(ruby: ruby, rbs: rbs)
+        end
+      end
+
+      def initialize(ruby:, rbs:)
+        raise ArgumentError, 'ruby must have ".rb" extension' unless ruby.type == :ruby
+        raise ArgumentError, 'rbs must have ".rbs" extension' if !rbs.nil? && rbs.type != :rbs
+
+        @ruby = ruby
+        @rbs = rbs
+      end
+
+      def to_s
+        rbs.nil? ? ruby.to_s : "#{ruby}\n#{rbs}"
+      end
+
+      attr_reader :ruby, :rbs
+    end
+  end
+end

data/lib/rbs_goose/io.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require_relative 'io/example'
+require_relative 'io/example_group'
+require_relative 'io/file'
+require_relative 'io/target_group'
+require_relative 'io/typed_ruby'
+
+module RbsGoose
+  module IO
+  end
+end

data/lib/rbs_goose/templates/default_template.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'langchain'
+
+module RbsGoose
+  module Templates
+    class DefaultTemplate
+      def initialize(instruction:, example_groups:)
+        @template = Langchain::Prompt::FewShotPromptTemplate.new(
+          prefix: instruction,
+          suffix: "#{input_template_string}\n",
+          example_prompt: example_prompt,
+          examples: example_groups.map { transform_example_group(_1) },
+          input_variables: %w[typed_ruby_list]
+        )
+      end
+
+      def format(typed_ruby_list)
+        template.format(typed_ruby_list: typed_ruby_list.join("\n"))
+      end
+
+      def parse_result(result)
+        result.scan(/```.+?```/m).map { IO::File.from_markdown(_1) }
+      end
+
+      private
+
+      attr_reader :template
+
+      def example_prompt
+        Langchain::Prompt::PromptTemplate.new(
+          template: "#{input_template_string}\n{refined_rbs_list}",
+          input_variables: %w[typed_ruby_list refined_rbs_list]
+        )
+      end
+
+      def input_template_string
+        "========Input========\n{typed_ruby_list}\n\n========Output========"
+      end
+
+      def transform_example_group(example_group)
+        {
+          typed_ruby_list: example_group.map(&:typed_ruby).join("\n"),
+          refined_rbs_list: example_group.map(&:refined_rbs).join("\n")
+        }
+      end
+    end
+  end
+end

data/lib/rbs_goose/templates.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require_relative 'templates/default_template'
+
+module RbsGoose
+  module Templates
+  end
+end

data/lib/rbs_goose/type_inferrer.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require_relative 'templates'
+
+module RbsGoose
+  class TypeInferrer
+    def infer(target_group)
+      template = RbsGoose.configuration.template
+      result = RbsGoose.llm.complete(prompt: template.format(target_group)).completion
+      template.parse_result(result)
+    end
+  end
+end

data/lib/rbs_goose/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module RbsGoose
+  VERSION = '0.1.0'
+end

data/lib/rbs_goose.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require_relative 'rbs_goose/configuration'
+require_relative 'rbs_goose/io'
+require_relative 'rbs_goose/type_inferrer'
+require_relative 'rbs_goose/version'
+
+require 'forwardable'
+
+module RbsGoose
+  class Error < StandardError; end
+
+  class << self
+    extend Forwardable
+
+    def configure(&block)
+      @configuration = Configuration.new(&block)
+    end
+
+    def reset_configuration
+      @configuration = nil
+    end
+
+    def run(code_dir: 'lib', sig_dir: 'sig', base_path: ::Dir.pwd)
+      puts "Run RbsGoose.(Code Directory: #{code_dir}, Signature Directory: #{sig_dir})"
+      target_group = RbsGoose::IO::TargetGroup.load_from(base_path, code_dir: code_dir, sig_dir: sig_dir)
+      RbsGoose::TypeInferrer.new.infer(target_group).each do |refined_rbs|
+        puts "write refined rbs to #{refined_rbs.path}\n"
+        refined_rbs.write
+        puts "done.\n\n"
+      end
+    end
+
+    attr_reader :configuration
+
+    def_delegators :configuration, :llm, :instruction, :example_groups
+  end
+end