Prototype: Add annotations command

README.md

@@ -673,6 +673,8 @@ check_shims:
   dsl_rbi_dir: sorbet/rbi/dsl
   shim_rbi_dir: sorbet/rbi/shims
   payload: true
+annotations:
+  repo_uri: https://raw.githubusercontent.com/Shopify/rbi-central/main
 ```
 <!-- END_CONFIG_TEMPLATE -->
 

lib/tapioca.rb

@@ -35,12 +35,17 @@ class Error < StandardError; end
   DEFAULT_GEM_DIR = T.let("#{DEFAULT_RBI_DIR}/gems", String)
   DEFAULT_SHIM_DIR = T.let("#{DEFAULT_RBI_DIR}/shims", String)
   DEFAULT_TODO_FILE = T.let("#{DEFAULT_RBI_DIR}/todo.rbi", String)
+  DEFAULT_ANNOTATIONS_DIR = T.let("#{DEFAULT_RBI_DIR}/annotations", String)
 
   DEFAULT_OVERRIDES = T.let({
     # ActiveSupport overrides some core methods with different signatures
     # so we generate a typed: false RBI for it to suppress errors
     "activesupport" => "false",
   }.freeze, T::Hash[String, String])
+
+  CENTRAL_REPO_ROOT_URI = "https://raw.githubusercontent.com/Shopify/rbi-central/main"
+  CENTRAL_REPO_INDEX_PATH = "index.json"
+  CENTRAL_REPO_ANNOTATIONS_DIR = "rbi/annotations"
 end
 
 require "tapioca/version"

lib/tapioca/cli.rb

@@ -301,6 +301,13 @@ def check_shims
       exit(0)
     end
 
+    desc "annotations", "Pull gem annotations from a central RBI repository"
+    option :repo_uri, type: :string, desc: "Repository URI to pull annotations from", default: CENTRAL_REPO_ROOT_URI
+    def annotations
+      command = Commands::Annotations.new(central_repo_root_uri: options[:repo_uri])
+      command.execute
+    end
+
     map ["--version", "-v"] => :__print_version
 
     desc "--version, -v", "show version"

lib/tapioca/commands.rb

@@ -4,6 +4,7 @@
 module Tapioca
   module Commands
     autoload :Command, "tapioca/commands/command"
+    autoload :Annotations, "tapioca/commands/annotations"
     autoload :Dsl, "tapioca/commands/dsl"
     autoload :Init, "tapioca/commands/init"
     autoload :Gem, "tapioca/commands/gem"

lib/tapioca/commands/annotations.rb

@@ -0,0 +1,154 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "net/http"
+
+module Tapioca
+  module Commands
+    class Annotations < Command
+      extend T::Sig
+
+      sig do
+        params(
+          central_repo_root_uri: String,
+          central_repo_index_path: String
+        ).void
+      end
+      def initialize(central_repo_root_uri:, central_repo_index_path: CENTRAL_REPO_INDEX_PATH)
+        super()
+        @central_repo_root_uri = central_repo_root_uri
+        @index = T.let(fetch_index, RepoIndex)
+      end
+
+      sig { override.void }
+      def execute
+        project_gems = list_gemfile_gems
+        remove_expired_annotations(project_gems)
+        fetch_annotations(project_gems)
+      end
+
+      private
+
+      sig { returns(T::Array[String]) }
+      def list_gemfile_gems
+        say("Listing gems from Gemfile.lock... ", [:blue, :bold])
+        gemfile = Bundler.read_file("Gemfile.lock")
+        parser = Bundler::LockfileParser.new(gemfile)
+        gem_names = parser.specs.map(&:name)
+        say("Done", :green)
+        gem_names
+      end
+
+      sig { params(project_gems: T::Array[String]).void }
+      def remove_expired_annotations(project_gems)
+        say("Removing annotations for gems that have been removed... ", [:blue, :bold])
+
+        annotations = Pathname.glob("#{DEFAULT_ANNOTATIONS_DIR}/*.rbi").map { |f| f.basename(".*").to_s }
+        expired = annotations - project_gems
+
+        if expired.empty?
+          say(" Nothing to do")
+          return
+        end
+
+        say("\n")
+        expired.each do |gem_name|
+          say("\n")
+          path = "#{DEFAULT_ANNOTATIONS_DIR}/#{gem_name}.rbi"
+          remove_file(path)
+        end
+        say("\nDone\n\n", :green)
+      end
+
+      sig { returns(RepoIndex) }
+      def fetch_index
+        say("Retrieving index from central repository... ", [:blue, :bold])
+        content = fetch_file(CENTRAL_REPO_INDEX_PATH)
+        exit(1) unless content
+
+        index = RepoIndex.from_json(content)
+        say("Done", :green)
+        index
+      end
+
+      sig { params(gem_names: T::Array[String]).returns(T::Array[String]) }
+      def fetch_annotations(gem_names)
+        say("Fetching gem annotations from central repository... ", [:blue, :bold])
+        fetchable_gems = gem_names.select { |gem_name| @index.has_gem?(gem_name) }
+
+        if fetchable_gems.empty?
+          say(" Nothing to do")
+          exit(0)
+        end
+
+        say("\n")
+        fetched_gems = fetchable_gems.select { |name| fetch_annotation(name) }
+        say("\nDone", :green)
+        fetched_gems
+      end
+
+      sig { params(gem_name: String).void }
+      def fetch_annotation(gem_name)
+        content = fetch_file("#{CENTRAL_REPO_ANNOTATIONS_DIR}/#{gem_name}.rbi")
+        return unless content
+
+        content = add_header(gem_name, content)
+
+        dir = DEFAULT_ANNOTATIONS_DIR
+        FileUtils.mkdir_p(dir)
+        say("\n  Fetched #{set_color(gem_name, :yellow, :bold)}", :green)
+        create_file("#{dir}/#{gem_name}.rbi", content)
+      end
+
+      sig { params(path: String).returns(T.nilable(String)) }
+      def fetch_file(path)
+        if @central_repo_root_uri.start_with?(%r{https?://})
+          fetch_http_file(path)
+        else
+          fetch_local_file(path)
+        end
+      end
+
+      sig { params(path: String).returns(T.nilable(String)) }
+      def fetch_local_file(path)
+        File.read("#{@central_repo_root_uri}/#{path}")
+      rescue => e
+        say_error("\nCan't fetch file `#{path}` (#{e.message})", :bold, :red)
+        nil
+      end
+
+      sig { params(path: String).returns(T.nilable(String)) }
+      def fetch_http_file(path)
+        uri = URI("#{@central_repo_root_uri}/#{path}")
+        response = Net::HTTP.get_response(uri)
+        case response
+        when Net::HTTPSuccess
+          response.body
+        else
+          say_error("\nCan't fetch file `#{path}` from #{@central_repo_root_uri} (#{response.class})", :bold, :red)
+          nil
+        end
+      rescue SocketError, Errno::ECONNREFUSED => e
+        say_error("\nCan't fetch file `#{path}` from #{@central_repo_root_uri} (#{e.message})", :bold, :red)
+        nil
+      end
+
+      sig { params(name: String, content: String).returns(String) }
+      def add_header(name, content)
+        header = <<~COMMENT
+          # DO NOT EDIT MANUALLY
+          # This file was pulled from #{@central_repo_root_uri}.
+          # Please run `#{default_command(:annotations)}` to update it.
+        COMMENT
+
+        contents = content.split("\n")
+        if contents[0]&.start_with?("# typed:") && contents[1]&.empty?
+          contents.insert(2, header).join("\n")
+        else
+          say_error("Couldn't insert file header for content: #{content} due to unexpected file format")
+          content
+        end
+      end
+    end
+  end
+end

lib/tapioca/internal.rb

@@ -16,6 +16,7 @@
 require "tapioca/helpers/signatures_helper"
 require "tapioca/helpers/rbi_helper"
 require "tapioca/helpers/shims_helper"
+require "tapioca/repo_index"
 require "tapioca/commands"
 require "tapioca/cli"
 require "tapioca/gemfile"

lib/tapioca/repo_index.rb

@@ -0,0 +1,41 @@
+# typed: strict
+# frozen_string_literal: true
+
+module Tapioca
+  class RepoIndex
+    extend T::Sig
+    extend T::Generic
+
+    sig { params(json: String).returns(RepoIndex) }
+    def self.from_json(json)
+      RepoIndex.from_hash(JSON.parse(json))
+    end
+
+    sig { params(hash: T::Hash[String, T::Hash[T.untyped, T.untyped]]).returns(RepoIndex) }
+    def self.from_hash(hash)
+      hash.each_with_object(RepoIndex.new) do |(name, _), index|
+        index << name
+      end
+    end
+
+    sig { void }
+    def initialize
+      @entries = T.let(Set.new, T::Set[String])
+    end
+
+    sig { params(gem_name: String).void }
+    def <<(gem_name)
+      @entries.add(gem_name)
+    end
+
+    sig { returns(T::Enumerable[String]) }
+    def gems
+      @entries.sort
+    end
+
+    sig { params(gem_name: String).returns(T::Boolean) }
+    def has_gem?(gem_name)
+      @entries.include?(gem_name)
+    end
+  end
+end

spec/tapioca/cli/annotations_spec.rb

@@ -0,0 +1,126 @@
+# typed: true
+# frozen_string_literal: true
+
+require "spec_helper"
+
+module Tapioca
+  class AnnotationsTest < SpecWithProject
+    describe "cli::annotations" do
+      before(:all) do
+        @project.bundle_install
+      end
+
+      after do
+        @project.remove("sorbet/rbi/annotations")
+      end
+
+      it "does nothing if the repo is empty" do
+        repo = create_repo({})
+
+        result = @project.tapioca("annotations --repo-uri #{repo.path}")
+
+        assert_equal(<<~OUT, result.out)
+          Retrieving index from central repository... Done
+          Listing gems from Gemfile.lock... Done
+          Removing annotations for gems that have been removed...  Nothing to do
+          Fetching gem annotations from central repository...  Nothing to do
+        OUT
+
+        assert_success_status(result)
+
+        repo.destroy
+      end
+
+      it "removes local annotations if they do not appear in the Gemfile.lock" do
+        repo = create_repo({})
+
+        @project.write("sorbet/rbi/annotations/rbi.rbi", "# typed: true")
+        @project.write("sorbet/rbi/annotations/bar.rbi", "# typed: true")
+        @project.write("sorbet/rbi/annotations/foo.rbi", "# typed: true")
+
+        result = @project.tapioca("annotations --repo-uri #{repo.path}")
+
+        assert_includes(result.out, "remove  sorbet/rbi/annotations/bar.rbi")
+        assert_includes(result.out, "remove  sorbet/rbi/annotations/foo.rbi")
+        refute_includes(result.out, "remove  sorbet/rbi/annotations/rbi.rbi")
+
+        assert_success_status(result)
+
+        repo.destroy
+      end
+
+      it "gets annotations from the central repo" do
+        repo = create_repo({
+          rbi: <<~RBI,
+            # typed: true
+
+            class AnnotationForRBI; end
+          RBI
+          spoom: <<~RBI,
+            # typed: strict
+
+            class AnnotationForSpoom; end
+          RBI
+          foo: <<~RBI,
+            # typed: false
+
+            class AnnotationForFoo; end
+          RBI
+        })
+
+        result = @project.tapioca("annotations --repo-uri #{repo.path}")
+
+        assert_includes(result.out, "create  sorbet/rbi/annotations/rbi.rbi")
+        assert_includes(result.out, "create  sorbet/rbi/annotations/spoom.rbi")
+        refute_includes(result.out, "create  sorbet/rbi/annotations/foo.rbi")
+
+        assert_project_annotation_equal(repo, "sorbet/rbi/annotations/rbi.rbi", <<~RBI)
+          # typed: true
+
+          # DO NOT EDIT MANUALLY
+          # This file was pulled from $REPO_PATH.
+          # Please run `bin/tapioca annotations` to update it.
+
+          class AnnotationForRBI; end
+        RBI
+
+        assert_project_annotation_equal(repo, "sorbet/rbi/annotations/spoom.rbi", <<~RBI)
+          # typed: strict
+
+          # DO NOT EDIT MANUALLY
+          # This file was pulled from $REPO_PATH.
+          # Please run `bin/tapioca annotations` to update it.
+
+          class AnnotationForSpoom; end
+        RBI
+
+        refute_project_file_exist("sorbet/rbi/annotations/foo.rbi")
+        assert_success_status(result)
+
+        repo.destroy
+      end
+    end
+
+    private
+
+    sig { params(annotations: T::Hash[String, String]).returns(MockDir) }
+    def create_repo(annotations)
+      repo = MockDir.new("#{@project.path}/repo")
+      index = {}
+
+      annotations.each do |gem_name, rbi_content|
+        index[gem_name] = {}
+        repo.write("rbi/annotations/#{gem_name}.rbi", rbi_content)
+      end
+
+      repo.write("index.json", index.to_json)
+      repo
+    end
+
+    sig { params(repo: MockDir, path: String, content: String).void }
+    def assert_project_annotation_equal(repo, path, content)
+      rbi_annotation = @project.read(path)
+      assert_equal(content.strip, rbi_annotation.gsub(repo.path, "$REPO_PATH").strip)
+    end
+  end
+end