# typed: strict # frozen_string_literal: true module RBI module Rewriters class Annotate < Visitor extend T::Sig sig { params(annotation: String, annotate_scopes: T::Boolean, annotate_properties: T::Boolean).void } def initialize(annotation, annotate_scopes: false, annotate_properties: false) super() @annotation = annotation @annotate_scopes = annotate_scopes @annotate_properties = annotate_properties end sig { override.params(node: T.nilable(Node)).void } def visit(node) case node when Scope annotate_node(node) if @annotate_scopes || root?(node) when Const, Attr, Method, TStructField, TypeMember annotate_node(node) if @annotate_properties end visit_all(node.nodes) if node.is_a?(Tree) end private sig { params(node: NodeWithComments).void } def annotate_node(node) return if node.annotations.one?(@annotation) node.comments << Comment.new("@#{@annotation}") end sig { params(node: Node).returns(T::Boolean) } def root?(node) parent = node.parent_tree parent.is_a?(Tree) && parent.parent_tree.nil? end end end class Tree extend T::Sig sig { params(annotation: String, annotate_scopes: T::Boolean, annotate_properties: T::Boolean).void } def annotate!(annotation, annotate_scopes: false, annotate_properties: false) visitor = Rewriters::Annotate.new( annotation, annotate_scopes: annotate_scopes, annotate_properties: annotate_properties ) visitor.visit(self) end end end