Source Code for Module rhn.SmartIO

 1  # 
 2  # Smart IO class 
 3  # 
 4  # Copyright (c) 2002--2016 Red Hat, Inc. 
 5  # 
 6  # Author: Mihai Ibanescu <misa@redhat.com> 
 7   
 8  """ 
 9  This module implements the SmartIO class 
10  """ 
11   
12  import os 
13  try: # python2 
14      from cStringIO import StringIO 
15  except ImportError: # python3 
16      from io import BytesIO as StringIO 
17   
18 -class SmartIO: 
19      """ 
20      The SmartIO class allows one to put a cap on the memory consumption. 
21      StringIO objects are very fast, because they are stored in memory, but 
22      if they are too big the memory footprint becomes noticeable. 
23      The write method of a SmartIO determines if the data that is to be added 
24      to the (initially) StrintIO object does not exceed a certain threshold; if 
25      it does, it switches the storage to a temporary disk file 
26      """ 
27 -    def __init__(self, max_mem_size=16384, force_mem=0): 
28          self._max_mem_size = max_mem_size 
29          self._io = StringIO() 
30          # self._fixed is a flag to show if we're supposed to consider moving 
31          # the StringIO object into a tempfile 
32          # Invariant: if self._fixed == 0, we have a StringIO (if self._fixed 
33          # is 1 and force_mem was 0, then we have a file) 
34          if force_mem: 
35              self._fixed = 1 
36          else: 
37              self._fixed = 0 
38   
39 -    def set_max_mem_size(self, max_mem_size): 
40          self._max_mem_size = max_mem_size 
41   
42 -    def get_max_mem_size(self): 
43          return self._max_mem_size 
44   
45 -    def write(self, data): 
46          if not self._fixed: 
47              # let's consider moving it to a file 
48              if len(data) + self._io.tell() > self._max_mem_size: 
49                  # We'll overflow, change to a tempfile 
50                  tmpfile = _tempfile() 
51                  tmpfile.write(self._io.getvalue()) 
52                  self._fixed = 1 
53                  self._io = tmpfile 
54   
55          self._io.write(data) 
56   
57 -    def __getattr__(self, name): 
58          return getattr(self._io, name) 
59   
60  # Creates a temporary file and passes back its file descriptor 
61 -def _tempfile(): 
62      import tempfile 
63      (fd, fname) = tempfile.mkstemp(prefix="_rhn_transports-%d-" \ 
64                                     % os.getpid()) 
65      # tempfile, unlink it 
66      os.unlink(fname) 
67      return os.fdopen(fd, "wb+") 
68