lib/pelusa/analyzer.rb in pelusa-0.2.2 vs lib/pelusa/analyzer.rb in pelusa-0.2.3

@@ -17,13 +17,14 @@
     # Returns a Report of all the classes.
     def analyze(ast)
       reports = extract_classes(ast).map do |klass|
         class_analyzer = ClassAnalyzer.new(klass)
         class_name     = class_analyzer.class_name
+        type           = class_analyzer.type
         analysis       = class_analyzer.analyze(@lints)
 
-        Report.new(class_name, analysis)
+        Report.new(class_name, type, analysis)
       end
       @reporter.reports = reports
       @reporter
     end
 
@@ -36,13 +37,19 @@
     # ast - The abstract syntax tree to extract the classes from.
     #
     # Returns an Array of Class nodes.
     def extract_classes(ast)
       classes = []
-      class_iterator = Iterator.new do |node|
-        classes << node if node.is_a?(Rubinius::AST::Class)
+      if ast.is_a?(Rubinius::AST::Class) || ast.is_a?(Rubinius::AST::Module)
+        classes << ast
       end
-      Array(ast).each(&class_iterator)
+
+      ast.walk do |continue, node|
+        if node.is_a?(Rubinius::AST::Class) || node.is_a?(Rubinius::AST::Module)
+          classes << node
+        end
+        true
+      end
       classes
     end
   end
 end