1. Gregory Petukhov
  2. grab

Commits

Gregory Petukhov  committed 33c8000

Add new option: body_storage_filename. Write tests for options: body_inmemory, body_storage_dir, body_storage_filename

  • Participants
  • Parent commits 7ead7db
  • Branches default

Comments (0)

Files changed (3)

File grab/base.py

View file
         body_maxsize = None,
         body_inmemory = True,
         body_storage_dir = None,
+        body_storage_filename = None,
 
         # Content compression
         encoding = 'gzip',
         # It's important to delete old POST data after request is performed.
         # If POST data is not cleared then next request will try to use them again!
 
-        self.config['post'] = None
-        self.config['multipart_post'] = None
-        self.config['method'] = None
+        self.reset_temporary_options()
 
         if prepare_response_func:
             self.response = prepare_response_func(self.transport, self)
 
         return None
 
+    def reset_temporary_options(self):
+        self.config['post'] = None
+        self.config['multipart_post'] = None
+        self.config['method'] = None
+        self.config['body_storage_filename'] = None
+
     def save_failed_dump(self):
         """
         Save dump of failed request for debugging.

File grab/transport/curl.py

View file
     from urllib.parse import urlsplit, urlunsplit
 import pycurl
 import tempfile
+import os.path
 
 from ..base import UploadContent, UploadFile
 from .. import error
     def __init__(self):
         self.curl = pycurl.Curl()
 
-    def setup_body_file(self, storage_dir):
-        handle, path = tempfile.mkstemp(dir=storage_dir)
+    def setup_body_file(self, storage_dir, storage_filename):
+        if storage_filename is None:
+            handle, path = tempfile.mkstemp(dir=storage_dir)
+        else:
+            path = os.path.join(storage_dir, storage_filename)
         self.body_file = open(path, 'wb')
         self.body_path = path
 
             self.curl.setopt(pycurl.WRITEFUNCTION, self.body_processor)
         else:
             if not grab.config['body_storage_dir']:
-                raise error.GrabMisuseError('Option body_storage_di is not defined')
-            self.setup_body_file(grab.config['body_storage_dir'])
+                raise error.GrabMisuseError('Option body_storage_dir is not defined')
+            self.setup_body_file(grab.config['body_storage_dir'],
+                                 grab.config['body_storage_filename'])
             self.curl.setopt(pycurl.WRITEFUNCTION, self.body_processor)
 
         if grab.config['verbose_logging']:

File test/base_interface.py

View file
 # coding: utf-8
 from unittest import TestCase
+import os
 
 from grab import Grab, GrabMisuseError
 from util import (FakeServerThread, BASE_URL, RESPONSE, REQUEST,
                   RESPONSE_ONCE, ignore_transport, GRAB_TRANSPORT,
-                  ignore_transport, only_transport)
+                  ignore_transport, only_transport,
+                  TMP_DIR)
 from grab.extension import register_extensions
 
 
         #SimpleExtension.get_data()['counter'] = 0
         #g = VeryCustomGrab()
         #self.assertEqual(SimpleExtension.get_data()['counter'], 2)
+
+    def test_body_inmemory(self):
+        g = Grab()
+        g.setup(body_inmemory=False)
+        self.assertRaises(GrabMisuseError, lambda: g.go(BASE_URL))
+
+        RESPONSE['get'] = 'foo'
+        g = Grab()
+        g.setup(body_inmemory=False)
+        g.setup(body_storage_dir=TMP_DIR)
+        g.go(BASE_URL)
+        self.assertTrue(os.path.exists(g.response.body_path))
+        self.assertTrue(TMP_DIR in g.response.body_path)
+        self.assertEqual('foo', open(g.response.body_path).read())
+        old_path = g.response.body_path
+
+        g.go(BASE_URL)
+        self.assertTrue(old_path != g.response.body_path)
+
+        RESPONSE['get'] = 'foo'
+        g = Grab()
+        g.setup(body_inmemory=False)
+        g.setup(body_storage_dir=TMP_DIR)
+        g.setup(body_storage_filename='musik.mp3')
+        g.go(BASE_URL)
+        self.assertTrue(os.path.exists(g.response.body_path))
+        self.assertTrue(TMP_DIR in g.response.body_path)
+        self.assertEqual('foo', open(g.response.body_path).read())
+        self.assertEqual(os.path.join(TMP_DIR, 'musik.mp3'), g.response.body_path)