openerp/tools/osutil.py

   1 # -*- coding: utf-8 -*-
   2 ##############################################################################
   3 #
   4 #    OpenERP, Open Source Management Solution
   5 #    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>).
   6 #
   7 #    This program is free software: you can redistribute it and/or modify
   8 #    it under the terms of the GNU Affero General Public License as
   9 #    published by the Free Software Foundation, either version 3 of the
  10 #    License, or (at your option) any later version.
  11 #
  12 #    This program is distributed in the hope that it will be useful,
  13 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 #    GNU Affero General Public License for more details.
  16 #
  17 #    You should have received a copy of the GNU Affero General Public License
  18 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
  19 #
  20 ##############################################################################
  21
  22 """
  23 Some functions related to the os and os.path module
  24 """
  25
  26 from contextlib import contextmanager
  27 import os
  28 from os.path import join as opj
  29 import shutil
  30 import tempfile
  31 import zipfile
  32
  33 if os.name == 'nt':
  34     import ctypes
  35     import win32service as ws
  36     import win32serviceutil as wsu
  37
  38
  39 def listdir(dir, recursive=False):
  40     """Allow to recursively get the file listing"""
  41     dir = os.path.normpath(dir)
  42     if not recursive:
  43         return os.listdir(dir)
  44
  45     res = []
  46     for root, dirs, files in walksymlinks(dir):
  47         root = root[len(dir)+1:]
  48         res.extend([opj(root, f) for f in files])
  49     return res
  50
  51 def walksymlinks(top, topdown=True, onerror=None):
  52     """
  53     same as os.walk but follow symlinks
  54     attention: all symlinks are walked before all normals directories
  55     """
  56     for dirpath, dirnames, filenames in os.walk(top, topdown, onerror):
  57         if topdown:
  58             yield dirpath, dirnames, filenames
  59
  60         symlinks = filter(lambda dirname: os.path.islink(os.path.join(dirpath, dirname)), dirnames)
  61         for s in symlinks:
  62             for x in walksymlinks(os.path.join(dirpath, s), topdown, onerror):
  63                 yield x
  64
  65         if not topdown:
  66             yield dirpath, dirnames, filenames
  67
  68 @contextmanager
  69 def tempdir():
  70     tmpdir = tempfile.mkdtemp()
  71     try:
  72         yield tmpdir
  73     finally:
  74         shutil.rmtree(tmpdir)
  75
  76 def zip_dir(path, stream, include_dir=True):      # TODO add ignore list
  77     path = os.path.normpath(path)
  78     len_prefix = len(os.path.dirname(path)) if include_dir else len(path)
  79     if len_prefix:
  80         len_prefix += 1
  81
  82     with zipfile.ZipFile(stream, 'w', compression=zipfile.ZIP_DEFLATED, allowZip64=True) as zipf:
  83         for dirpath, dirnames, filenames in os.walk(path):
  84             for fname in filenames:
  85                 bname, ext = os.path.splitext(fname)
  86                 ext = ext or bname
  87                 if ext not in ['.pyc', '.pyo', '.swp', '.DS_Store']:
  88                     path = os.path.normpath(os.path.join(dirpath, fname))
  89                     if os.path.isfile(path):
  90                         zipf.write(path, path[len_prefix:])
  91
  92
  93 if os.name != 'nt':
  94     getppid = os.getppid
  95     is_running_as_nt_service = lambda: False
  96 else:
  97     # based on http://mail.python.org/pipermail/python-win32/2007-June/006174.html
  98     _TH32CS_SNAPPROCESS = 0x00000002
  99     class _PROCESSENTRY32(ctypes.Structure):
 100         _fields_ = [("dwSize", ctypes.c_ulong),
 101                     ("cntUsage", ctypes.c_ulong),
 102                     ("th32ProcessID", ctypes.c_ulong),
 103                     ("th32DefaultHeapID", ctypes.c_ulong),
 104                     ("th32ModuleID", ctypes.c_ulong),
 105                     ("cntThreads", ctypes.c_ulong),
 106                     ("th32ParentProcessID", ctypes.c_ulong),
 107                     ("pcPriClassBase", ctypes.c_ulong),
 108                     ("dwFlags", ctypes.c_ulong),
 109                     ("szExeFile", ctypes.c_char * 260)]
 110
 111     def getppid():
 112         CreateToolhelp32Snapshot = ctypes.windll.kernel32.CreateToolhelp32Snapshot
 113         Process32First = ctypes.windll.kernel32.Process32First
 114         Process32Next = ctypes.windll.kernel32.Process32Next
 115         CloseHandle = ctypes.windll.kernel32.CloseHandle
 116         hProcessSnap = CreateToolhelp32Snapshot(_TH32CS_SNAPPROCESS, 0)
 117         current_pid = os.getpid()
 118         try:
 119             pe32 = _PROCESSENTRY32()
 120             pe32.dwSize = ctypes.sizeof(_PROCESSENTRY32)
 121             if not Process32First(hProcessSnap, ctypes.byref(pe32)):
 122                 raise OSError('Failed getting first process.')
 123             while True:
 124                 if pe32.th32ProcessID == current_pid:
 125                     return pe32.th32ParentProcessID
 126                 if not Process32Next(hProcessSnap, ctypes.byref(pe32)):
 127                     return None
 128         finally:
 129             CloseHandle(hProcessSnap)
 130
 131     from contextlib import contextmanager
 132     from openerp.release import nt_service_name
 133
 134     def is_running_as_nt_service():
 135         @contextmanager
 136         def close_srv(srv):
 137             try:
 138                 yield srv
 139             finally:
 140                 ws.CloseServiceHandle(srv)
 141
 142         try:
 143             with close_srv(ws.OpenSCManager(None, None, ws.SC_MANAGER_ALL_ACCESS)) as hscm:
 144                 with close_srv(wsu.SmartOpenService(hscm, nt_service_name, ws.SERVICE_ALL_ACCESS)) as hs:
 145                     info = ws.QueryServiceStatusEx(hs)
 146                     return info['ProcessId'] == getppid()
 147         except Exception:
 148             return False
 149
 150 if __name__ == '__main__':
 151     from pprint import pprint as pp
 152     pp(listdir('../report', True))
 153
 154 # vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4: