Move all of the 'src' code under the more-standard 'lib'.

[dead/whatever-dl.git] / lib / whatever-dl / uri_utilities.rb

diff --git a/lib/whatever-dl/uri_utilities.rb b/lib/whatever-dl/uri_utilities.rb
new file mode 100644 (file)
index 0000000..78966e8
--- /dev/null
+++ b/lib/whatever-dl/uri_utilities.rb
@@ -0,0 +1,85 @@
+#
+# Copyright Michael Orlitzky
+#
+# http://michael.orlitzky.com/
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# http://www.fsf.org/licensing/licenses/gpl.html
+#
+
+require 'open-uri'
+require 'vendor/ruby-progressbar/progressbar'
+
+# Just a couple of convenience methods for URIs.
+# These could be monkey-patched in, but with only
+# a few methods, that'd be asking (unnecessarily)
+# for trouble.
+class UriUtilities
+
+  def initialize()
+    @pbar = nil
+  end
+
+  def content_length_callback(content_length)
+    if content_length && (0 < content_length)
+      @pbar = ProgressBar.new("Download", content_length)
+      @pbar.instance_eval { @bar_mark = '=' }
+      @pbar.file_transfer_mode
+    end
+  end
+
+  def progress_callback(size)
+    @pbar.set(size) if @pbar
+  end
+
+
+  # Download the given URI object to <outfile_name>.
+  # Should use the progress_proc parameter to show
+  # a progress bar using the Ruby/ProgressBar library.
+  def download_with_progress_bar(uri, outfile_name, headers = {})
+    # We wrap the whole thing in a begin/rescue so that
+    # we can clean up afterwards in case of an error.
+    begin
+      File.open(outfile_name, 'wb') do |outfile|
+        # Convert our methods to Proc objects.
+        content_length_callback = method(:content_length_callback)
+        progress_callback = method(:progress_callback)
+
+        # So that they can be passed as arguments to open().
+        open_params = { :content_length_proc => content_length_callback,
+                        :progress_proc => progress_callback}
+
+        # Add the headers as additional parameters to the open() call.
+        open_params.merge!(headers)
+
+        uri.open(open_params) do |video_file|
+          outfile.write(video_file.read)
+        end
+      end
+
+      # Toss out an empty line to get rid of the progress bar.
+      # Normally, it would remain on the shell's "current" line.
+      puts ''
+
+    rescue Errno::EACCES => e
+      # Don't delete the file if it's unwritable.
+      raise(e)
+    rescue => e
+      # Here we get rid of the output file if there was an error.
+      # We test File.exists? first since the first line, the open()
+      # call, could theoretically fail.
+      File.delete(outfile_name) if File.exists?(outfile_name)
+      raise(e)
+    end
+  end
+
+end