Support A = Data.define(:a, :b, :c)

Author: marshall-lee

lib/yard/handlers/ruby/constant_handler.rb

@@ -9,6 +9,9 @@ class YARD::Handlers::Ruby::ConstantHandler < YARD::Handlers::Ruby::Base
     if statement[1].call? && statement[1][0][0] == s(:const, "Struct") &&
        statement[1][2] == s(:ident, "new")
       process_structclass(statement)
+    elsif statement[1].call? && statement[1][0][0] == s(:const, "Data") &&
+       statement[1][2] == s(:ident, "define")
+      process_dataclass(statement)
     elsif statement[0].type == :var_field && statement[0][0].type == :const
       process_constant(statement)
     elsif statement[0].type == :const_path_field
@@ -41,10 +44,24 @@ def process_structclass(statement)
     end
   end
 
-  # Extract the parameters from the Struct.new AST node, returning them as a list
+  def process_dataclass(statement)
+    lhs = statement[0]
+    if (lhs.type == :var_field && lhs[0].type == :const) || lhs.type == :const_path_field
+      klass = create_class(lhs.source, P(:Data))
+      extract_parameters(statement[1]).each do |member|
+        klass.attributes[:instance][member] = SymbolHash[:read => nil, :write => nil]
+        create_reader(klass, member)
+      end
+      parse_block(statement[1].block[1], :namespace => klass) unless statement[1].block.nil?
+    else
+      raise YARD::Parser::UndocumentableError, "Data assignment to #{lhs.source}"
+    end
+  end
+
+  # Extract the parameters from the Struct.new or Data.define AST node, returning them as a list
   # of strings
   #
-  # @param [MethodCallNode] superclass the AST node for the Struct.new call
+  # @param [MethodCallNode] superclass the AST node for the Struct.new or Data.define call
   # @return [Array<String>] the member names to generate methods for
   def extract_parameters(superclass)
     return [] unless superclass.parameters

spec/handlers/constant_handler_spec.rb

@@ -73,6 +73,48 @@
     expect(a3.tag(:return).types).to eq ["Symbol"]
   end
 
+  it "turns Const = Data.define(:sym) into class Const with attr reader :sym" do
+    obj = Registry.at("MyData")
+    expect(obj).to be_kind_of(CodeObjects::ClassObject)
+    expect(obj.superclass).to eq P(:Data)
+    attrs = obj.attributes[:instance]
+    [:a, :b, :c].each do |key|
+      expect(attrs).to have_key(key)
+      expect(attrs[key][:read]).not_to be nil
+      expect(attrs[key][:write]).to be nil
+    end
+  end
+
+  it "turns Const = Data.define into empty data class" do
+    obj = Registry.at("MyEmptyData")
+    expect(obj).to be_kind_of(CodeObjects::ClassObject)
+    expect(obj.superclass).to eq P(:Data)
+    expect(obj.attributes[:instance]).to be_empty
+  end
+
+  it "turns A::Const = Data.define(:sym) into class A::Const with attr reader :sym" do
+    obj = Registry.at("A::NestedCompactData")
+    expect(obj).to be_kind_of(CodeObjects::ClassObject)
+    attrs = obj.attributes[:instance]
+    [:b, :c].each do |key|
+      expect(attrs).to have_key(key)
+      expect(attrs[key][:read]).not_to be nil
+      expect(attrs[key][:write]).to be nil
+    end
+  end
+
+  it "documents block for Data.define if present" do
+    obj = Registry.at("MyDataWithMethods")
+    expect(obj).to be_kind_of(CodeObjects::ClassObject)
+    attrs = obj.attributes[:instance]
+    [:c, :d].each do |key|
+      expect(attrs).to have_key(key)
+      expect(attrs[key][:read]).not_to be nil
+      expect(attrs[key][:write]).to be nil
+    end
+    expect(obj.meths).to include(P("MyDataWithMethods#foo"))
+  end
+
   it "raises undocumentable error in 1.9 parser for Struct.new assignment to non-const" do
     undoc_error "nonconst = Struct.new"
   end unless LEGACY_PARSER

spec/handlers/examples/constant_handler_001.rb.txt

@@ -33,3 +33,10 @@ end
 #   Attribute defined with the new syntax
 #   @return [Symbol] something useful
 DocstringStruct = Struct.new(:bar, :baz, :new_syntax)
+
+MyData = Data.define(:a, :b, :c)
+MyEmptyData = Data.define
+A::NestedCompactData = Data.define(:b, :c)
+MyDataWithMethods = Data.define(:c, :d) do
+  def foo; end
+end