import os
import sys
import unittest
from multiprocessing.managers import SyncManager
from contextlib import contextmanager
__license__ = 'MIT'
__email__ = 'mat@miga.me'
#
#
#
class SharedDiskSpaceProctor(SyncManager):
   def __init__(self):
      SyncManager.__init__(self)
         self.register(	'DiskSpaceProctor',
               callable=DiskSpaceProctor,
               exposed=('_can_write', '_close_write', 'make_dirs'))
#
#
#
class DiskSpaceProctorException(Exception):
   def __init__(self,msg,**kwds):
      for k,v in kwds.items():
         setattr(self,k,v)
      self.msg = msg
   def __str__(self):
      return '%s - %s'%(self.msg, str(self.__dict__)) 

#
#
#
class DiskSpaceProctor(object):
   def __init__(self):
      self._wRLock = threading.RLock()
      self._currentWrites = {}
   #
   def _can_write(self, path, fz):
      '''
         returns True if there is available space and
         considers current write operations from other
         processes
      '''
      with self._wRLock:
         try:
            path = self.__chk_cwd(path) 
            s = os.statvfs(os.path.dirname(path))
            ca = s.f_frsize * s.f_bavail
            if self._currentWrites.has_key(path):
               raise DiskSpaceProctorException('Concurrent writes not allowed to same path%s'%path)
            wbt = 0
            for v in self._currentWrites.values():
               wbt += v
            if (fz + wbt) < ca:
               self._currentWrites[path]=fz
               return True
            raise DiskSpaceProctorException (
               'Raising Exception "%s" from %s.%s()'%(
                  'Not enough space for requested file size',
                  self.__class__.__name__, str(inspect.stack()[0][3]))
            )
         except DiskSpaceProctorException, de:
            raise de
         except Exception, e:
            raise Exception (
               'Raising Exception "%s" from %s.%s()'%(
                  e, self.__class__.__name__, str(inspect.stack()[0][3]))
            )
   #
   def _close_write(self, path):
      '''
         remove path and its filesz out of current writes addition
      '''
      with self._wRLock:
         path = self.__chk_cwd(path)
         self._currentWrites.pop(path)
   #
   def make_dirs(self, path, mode=0754, potFilesz=4096):
      '''
         create full path, consider potential file write sz
      '''
      if os.path.exists(path): return True
      if self._can_write(path, potFilesz):
         os.makedirs(path, mode=mode)
         self._close_write(path)
         return True
      return False
   #
   def __chk_cwd(self, path):
      return os.path.join(os.getcwd(),path) if os.path.dirname(path)=='' else os.path.dirname(path)
#
#
#  
class Test(unittest.TestCase):
   def test_diskspace_proctor(self):
      sdsp = SharedDiskSpaceProctor()
      sdsp.start()
      dsp = sdsp.DiskSpaceProctor()
      dspLock = sdsp.Lock()
      testFilename = 'diskspaceproctortest.dat'
      if os.path.exists(testFilename):
         os.remove(testFilename)
      doWrite = False
      with dspLock: 
         doWrite = dsp._can_write(testFilename, 4096):
      if doWrite:    
         f = open(testFilename, 'w')
         f.write('for testing only')
         f.close()
      self.assertTrue(os.path.exists(testFilename))   
#  
if __name__ == '__main__':
   unittest.main()