From f53cbed1bb07ce2d7b7c5b03632f41811f668430 Mon Sep 17 00:00:00 2001
From: Charlie Sharpsteen <source@sharpsteen.net>
Date: Wed, 21 May 2014 22:14:53 -0700
Subject: [PATCH] Add rudimentary comment extractor

With this patch, the `yardoc` face will parse a manifest file and then extract
any comments associated with each node, class and type definition.
---
 lib/puppet/face/yardoc.rb  |  9 ++++--
 lib/puppetx/yardoc/util.rb | 62 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)
 create mode 100644 lib/puppetx/yardoc/util.rb

diff --git a/lib/puppet/face/yardoc.rb b/lib/puppet/face/yardoc.rb
index bbb53e6..5884d17 100644
--- a/lib/puppet/face/yardoc.rb
+++ b/lib/puppet/face/yardoc.rb
@@ -1,11 +1,12 @@
 require 'puppet/face'
+require 'puppetx/yardoc/util'
 
 Puppet::Face.define(:yardoc, '0.0.1') do
 
   action(:yardoc) do
     default
 
-    when_invoked do |*args|
+    when_invoked do |manifest, options|
 
       unless Puppet.features.yard?
         raise RuntimeError, "The 'yard' gem must be installed in order to use this face."
@@ -18,9 +19,11 @@ Puppet::Face.define(:yardoc, '0.0.1') do
       end
 
       parser = Puppet::Pops::Parser::Parser.new()
-      parse_result = parser.parse_file(args[0])
+      parse_result = parser.parse_file(manifest)
 
-      return 0
+      commentor = Puppetx::Yardoc::Commentor.new()
+
+      return commentor.get_comments(parse_result)
     end
   end
 end
diff --git a/lib/puppetx/yardoc/util.rb b/lib/puppetx/yardoc/util.rb
new file mode 100644
index 0000000..ecdc9c2
--- /dev/null
+++ b/lib/puppetx/yardoc/util.rb
@@ -0,0 +1,62 @@
+require 'puppetx'
+require 'puppet/pops'
+
+# TODO: A catch-all module for now. Everything in here should eventually move
+# to a designated home.
+module Puppetx::Yardoc
+  # Loosly based on the TreeDumper classes in Pops::Model.
+  class Puppetx::Yardoc::Commentor
+    def initialize
+      @docstring_visitor = Puppet::Pops::Visitor.new(self,'docstring')
+    end
+
+    def get_comments(parse_result)
+      @docstring_visitor.visit(parse_result)
+    end
+
+    private
+
+    COMMENT_PATTERN = /^\s*#.*\n/
+
+    def comments(o)
+      @docstring_visitor.visit(o)
+    end
+
+    def docstring_Factory(o)
+      comments(o.current)
+    end
+
+    def docstring_Program(o)
+      @source_text = o.source_text.lines.to_a
+      @locator = o.locator
+
+      o.definitions.map{|d| comments(d)}
+    end
+
+    # Extract comments from "Definition" objects. That is: nodes definitions,
+    # type definitions and class definitions.
+    def docstring_Definition(o)
+      line = @locator.line_for_offset(o.offset)
+      comments_before(line)
+    end
+
+    def comments_before(line)
+      comments = []
+
+      # FIXME: Horribly inefficient. Multiple copies. Generator pattern would
+      # be much better.
+      @source_text.slice(0, line-1).reverse.each do |line|
+        if COMMENT_PATTERN.match(line)
+          comments.unshift line
+        else
+          # No comment found on this line. We must be done piecing together a
+          # comment block.
+          break
+        end
+      end
+
+      # Stick everything back together.
+      comments.join
+    end
+  end
+end