lib/annotator/model.rb in annotator-0.0.5 vs lib/annotator/model.rb in annotator-0.0.6

@@ -1,17 +1,22 @@
 module Annotator
 
   # Represents a single model file and associated class
   class Model
 
-    def initialize(filename)
+    def initialize(filename, base_path)
       @filename = filename
+      @base_path = base_path
       @blocks = Hash.new {[]}
     end
 
     # Model class
     def klass
-      @filename.split('app/models/').last.split(/\.rb$/).first.camelize.constantize rescue nil
+      begin 
+        @filename.split(@base_path).last.split(/\.rb$/).first.camelize.constantize rescue nil
+      rescue Exception
+        nil
+      end
     end
 
     # Split file into 3 blocks: before attributes, attributes block, and after
     # If there's no attributes block, content will be in :after part (for easier insertion)
     def parse