1. import os
  2. from stat import S_ISDIR
  3. from errno import ENOTDIR
  4. from functools import partial
  5.  
  6. O_DIRECTORY = getattr(os, "O_DIRECTORY", None)
  7.  
  8. AT_FUNCTIONS = set()
  9. for name in dir(os):
  10.     if name.endswith("at") and hasattr(getattr(os, name), "__call__"):
  11.         if hasattr(getattr(os, name[:-2], None), "__call__"):
  12.             AT_FUNCTIONS.add(name)
  13. del name
  14.  
  15.        
  16. class atcontext:
  17.     def __init__(self, directory, closefd=True, cloexec=False):
  18.         self.cloexec = cloexec
  19.         if isinstance(directory, str):
  20.             self.directory = directory
  21.             self.dirfd = None
  22.             self.closefd = True
  23.         elif isinstance(directory, int):
  24.             self.directory = None
  25.             self.dirfd = directory
  26.             self.closefd = closefd
  27.         else:
  28.             raise TypeError(directory)
  29.        
  30.     def _check(self):
  31.         if self.directory is None and self.dirfd is None:
  32.             raise ValueError("I/O operation on closed file.")
  33.  
  34.     def __enter__(self):
  35.         self._check()
  36.         if self.dirfd is None:
  37.             flags = os.O_RDONLY
  38.             if self.cloexec:
  39.                 flags |= os.O_CLOEXEC
  40.             if O_DIRECTORY is not None:
  41.                 flags |= O_DIRECTORY
  42.             elif not S_ISDIR(os.stat(self.directory).st_mode):
  43.                 # must check early as open blocks on fifos (DoS)
  44.                 # see Linux's man(2) open, O_DIRECTORY
  45.                 raise NotADirectoryError(ENOTDIR, os.strerror(ENOTDIR), self.directory)
  46.             self.dirfd = os.open(self.directory, flags)
  47.         return self
  48.    
  49.     def __exit__(self, typ, value, tb):
  50.         if self.closefd:
  51.             os.close(self.dirfd)
  52.         self.directory = self.dirfd = None
  53.    
  54.     def __getattr__(self, name):
  55.         self._check()
  56.         nameat = name + "at"
  57.         if nameat in AT_FUNCTIONS:
  58.             return partial(getattr(os, nameat), self.dirfd)
  59.         raise AttributeError(name)
  60.    
  61.     def __dir__(self):
  62.         content = super().__dir__() + list(AT_FUNCTIONS)
  63.         return sorted(content)
  64.        
  65. if __name__ == "__main__":
  66.     with atcontext("/etc") as at:
  67.         print(dir(at))
  68.         f = at.open("fstab", os.O_RDONLY)
  69.         print(os.read(f, 50))
  70.         os.close(f)
  71.    
  72.     with atcontext("/etc/fstab") as at:
  73.         pass