improvements

Author: DannyBen

lib/bashly/commands/generate.rb

@@ -35,7 +35,8 @@ def upgrade_libs
           content = File.read file
 
           if content =~ /\[@bashly-upgrade (.+)\]/
-            library_name, args = $1.split(' ', 2)
+            args = $1.split ' '
+            library_name = args.shift
             upgrade file, library_name, *args
           end
         end
@@ -55,7 +56,7 @@ def upgrade(existing_file, library_name, *args)
 
       def upgrade!(existing_file, library_name, *args)
         library = Bashly::Library.new library_name, *args
-        file = library.files.select { |f| f[:path] == existing_file }.first
+        file = library.find_file existing_file
 
         if file
           File.deep_write file[:path], file[:content]

lib/bashly/library.rb

@@ -22,7 +22,7 @@ def initialize(name, *args)
     end
 
     def files
-      if config.is_a? Symbol
+      if custom_handler
         custom_handler.files
 
       else
@@ -34,16 +34,21 @@ def files
     end
 
     def post_install_message
-      if config.is_a? Symbol
+      if custom_handler
         custom_handler.post_install_message
       else
         config['post_install_message']
       end
     end
 
+    def find_file(path)
+      files.select { |f| f[:path] == path }.first
+    end
+
   private
 
     def custom_handler
+      return nil unless config.is_a? Symbol
       @custom_handler ||= Bashly::Libraries.const_get(config).new(*args)
     end