data/python/FreeFOAM/util.py

   1 #-------------------------------------------------------------------------------
   2 #               ______                _     ____          __  __
   3 #              |  ____|             _| |_  / __ \   /\   |  \/  |
   4 #              | |__ _ __ ___  ___ /     \| |  | | /  \  | \  / |
   5 #              |  __| '__/ _ \/ _ ( (| |) ) |  | |/ /\ \ | |\/| |
   6 #              | |  | | |  __/  __/\_   _/| |__| / ____ \| |  | |
   7 #              |_|  |_|  \___|\___|  |_|   \____/_/    \_\_|  |_|
   8 #
   9 #                   FreeFOAM: The Cross-Platform CFD Toolkit
  10 #
  11 # Copyright (C) 2008-2012 Michael Wild <themiwi@users.sf.net>
  12 #                         Gerber van der Graaf <gerber_graaf@users.sf.net>
  13 #-------------------------------------------------------------------------------
  14 # License
  15 #   This file is part of FreeFOAM.
  16 #
  17 #   FreeFOAM is free software: you can redistribute it and/or modify it
  18 #   under the terms of the GNU General Public License as published by the
  19 #   Free Software Foundation, either version 3 of the License, or (at your
  20 #   option) any later version.
  21 #
  22 #   FreeFOAM is distributed in the hope that it will be useful, but WITHOUT
  23 #   ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  24 #   FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  25 #   for more details.
  26 #
  27 #   You should have received a copy of the GNU General Public License
  28 #   along with FreeFOAM.  If not, see <http://www.gnu.org/licenses/>.
  29 #
  30 # Description
  31 #   Useful Python functionality
  32 #
  33 #-------------------------------------------------------------------------------
  34
  35 """Classes and functions to manage FreeFOAM cases."""
  36
  37 from FreeFOAM.compat import *
  38
  39 import sys as _sys
  40 import os as _os
  41 import os.path as _op
  42 import re as _re
  43 import FreeFOAM as _f
  44
  45 class NoPolyMeshDirectoryError(Exception):
  46    """Thrown if no polyMesh directory could be found"""
  47    def __init__(self):
  48       Exception.__init__(self)
  49
  50    def __str__(self):
  51       return 'Failed to find a polyMesh directory'
  52
  53 class Tee:
  54    """A simple class which writes both to a file and standard output.
  55
  56    It only implements the following functions from `file`:
  57       * `Tee.write`
  58       * `Tee.writelines`
  59       * `Tee.flush`
  60       * `Tee.fileno`
  61
  62    """
  63    def __init__(self, fname, *args):
  64       """Initializes a new instance of `Tee`.
  65
  66       Parameters
  67       ----------
  68       fname     : name of the file to write to
  69       mode      : file mode (as in `open()`)
  70       buffering : buffering option (as in `open()`)
  71
  72       """
  73       self._file = open(fname, *args)
  74       self._stdout = _os.fdopen(_os.dup(_sys.stdout.fileno()), 'w')
  75
  76    def write(self, string):
  77       """Write `string` to the file and standard output"""
  78       self._file.write(string)
  79       self._stdout.write(string)
  80
  81    def writelines(self, strings):
  82       """Write `strings` to the file and standard output"""
  83       self._file.writelines(strings)
  84       self._stdout.writelines(strings)
  85
  86    def flush(self):
  87       """Flush the streams"""
  88       self._file.flush()
  89       self._stdout.flush()
  90
  91    def fileno(self):
  92       """Return the handle of the file stream"""
  93       return self._file.fileno()
  94
  95 def clone_case(old_case, new_case, verbose=True):
  96    """Clone a FreeFOAM case if the `new_case+'/system'` directory doesn't exist
  97
  98
  99    Parameters
 100    ----------
 101    old_case : The source case to be cloned.
 102    new_case : The destination case to be created.
 103
 104    """
 105    if not _op.exists(_op.join(new_case, 'system')):
 106       if verbose: echo('Cloning case %s'%old_case)
 107       for d in ['0', 'system', 'constant', 'chemkin']:
 108          if _op.isdir(_op.join(old_case, d)):
 109             copytree(
 110                   _op.join(old_case, d),
 111                   _op.join(new_case, d),
 112                   symlinks=True)
 113
 114 def remove_case(case='.', verbose=True):
 115    """Removes a case directory"""
 116    echo('Removing case %s'%case)
 117    rmtree(case)
 118
 119 def clean_case(case='.', verbose=True):
 120    """Clean a case by removing generated files and directories"""
 121    import glob
 122    import shutil
 123    case=_op.abspath(case)
 124    if verbose: echo('Cleaning case', case)
 125    patterns = ['[1-9]*', '-[1-9]*', 'log', 'log.*', 'log-*', 'logSummary.*',
 126          '.fxLock', '*.xml', 'ParaView*', 'paraFoam*', '*.OpenFOAM',
 127          'processor*', 'probes*', 'forces*', 'system/machines', 'VTK']
 128    for nZeros in range(8):
 129        timeDir = '0.'+nZeros*'0'+'[1-9]*'
 130        for p in ['', '-']:
 131           patterns.append(_op.join(case, p+timeDir))
 132    for p in ('allOwner cell face meshModifiers owner neighbour point edge'+
 133          'cellLevel pointLevel refinementHistory surfaceIndex sets').split():
 134       patterns.append(_op.join(case, 'constant', 'polyMesh', p+'*'))
 135    for p in 'cellToRegion cellLevel pointLevel'.split():
 136       patterns.append(_op.join(case, 'constant', p+'*'))
 137    for p in patterns:
 138       for dd in glob.glob(_op.join(case, p)):
 139          rmtree(dd)
 140
 141 def clean_samples(case='.'):
 142    """Remove sample directories"""
 143    for d in 'sets samples sampleSurfaces'.split():
 144       rmtree(d)
 145
 146 def rmtree(path, ignore_errors=False, onerror=None):
 147    """Recursively delete all files/directories under `path`
 148    (if `path` exists)"""
 149    import shutil
 150    import os
 151    path = _op.abspath(path)
 152    if _op.isdir(path):
 153       shutil.rmtree(path, ignore_errors, onerror)
 154    elif _op.isfile(path) or _op.islink(path):
 155       os.remove(path)
 156
 157 def copytree(src, dst, symlinks=True):
 158    """Recursively copies `src` to `dst`. If `dst` exists,
 159    it is removed first."""
 160    import shutil
 161    src = _op.abspath(src)
 162    dst = _op.abspath(dst)
 163    rmtree(dst)
 164    parent = _op.dirname(dst)
 165    if not _op.isdir(parent):
 166       _os.makedirs(parent)
 167    if _op.isdir(src):
 168       shutil.copytree(src, dst, symlinks)
 169    else:
 170       shutil.copy2(src, dst)
 171
 172 def rename(src, dst):
 173    """Forcebly renames `src` to `dst`. `dst` is removed first if it exists."""
 174    import shutil
 175    src = _op.normpath(src)
 176    dst = _op.normpath(dst)
 177    if _op.exists(dst):
 178       rmtree(dst)
 179    shutil.move(src, dst)
 180
 181
 182 def clear_polymesh(case=None, region=None):
 183    """Remove the contents of the constant/polyMesh directory as per the
 184    Foam::polyMesh::removeFiles() method.
 185
 186    Parameters
 187    ----------
 188    case   : Case directory from which to clear the mesh
 189    region : Mesh region to clear.
 190
 191    """
 192    meshDir = 'polyMesh'
 193    if region:
 194       meshDir = _op.join(region, meshDir)
 195    tmp = _op.join('constant', meshDir)
 196    if case:
 197       meshDir = _op.join(case, tmp)
 198    else:
 199       if _op.isdir(tmp):
 200          meshDir = tmp
 201       elif _op.isdir(meshDir):
 202          # probably already in 'constant'
 203          pass
 204       elif _op.basename(_os.getcwd()) == 'polyMesh':
 205          # probably already in 'polyMesh'
 206          meshDir = _os.getcwd()
 207
 208    if not _op.isdir(meshDir):
 209       raise NoPolyMeshDirectoryError()
 210
 211    for f in """points faces owner neighbour cells boundary pointZones faceZones
 212    cellZones meshModifiers parallelData sets cellLevel pointLevel
 213    refinementHistory surfaceIndex""".split():
 214       f = _op.join(meshDir, f)
 215       if _op.isfile(f):
 216          _os.remove(f)
 217
 218 def make_short_path(path, maxlen=40):
 219    """Abbreviate `path` to a maximum length of `maxlen` characters."""
 220    path = _op.abspath(path)
 221    first = 2-maxlen
 222    if len(path) > maxlen:
 223       idx = path.find(_os.sep, first)
 224       if idx < 0:
 225          idx = first
 226       path = '...'+path[idx:]
 227    return path
 228
 229 class TerminalController:
 230    """
 231    A class that can be used to portably generate formatted output to
 232    a terminal.
 233
 234    `TerminalController` defines a set of instance variables whose
 235    values are initialized to the control sequence necessary to
 236    perform a given action.  These can be simply included in normal
 237    output to the terminal:
 238
 239        >>> term = TerminalController()
 240        >>> print 'This is '+term.GREEN+'green'+term.NORMAL
 241
 242    Alternatively, the `render()` method can used, which replaces
 243    '${action}' with the string required to perform 'action':
 244
 245        >>> term = TerminalController()
 246        >>> print term.render('This is ${GREEN}green${NORMAL}')
 247
 248    If the terminal doesn't support a given action, then the value of
 249    the corresponding instance variable will be set to ''.  As a
 250    result, the above code will still work on terminals that do not
 251    support color, except that their output will not be colored.
 252    Also, this means that you can test whether the terminal supports a
 253    given action by simply testing the truth value of the
 254    corresponding instance variable:
 255
 256        >>> term = TerminalController()
 257        >>> if term.CLEAR_SCREEN:
 258        ...     print 'This terminal supports clearning the screen.'
 259
 260    Finally, if the width and height of the terminal are known, then
 261    they will be stored in the `COLS` and `LINES` attributes.
 262
 263    Copied from http://code.activestate.com/recipes/475116-using-terminfo-for-portable-color-output-cursor-co
 264    """
 265    # Cursor movement:
 266    BOL = ''             #: Move the cursor to the beginning of the line
 267    UP = ''              #: Move the cursor up one line
 268    DOWN = ''            #: Move the cursor down one line
 269    LEFT = ''            #: Move the cursor left one char
 270    RIGHT = ''           #: Move the cursor right one char
 271
 272    # Deletion:
 273    CLEAR_SCREEN = ''    #: Clear the screen and move to home position
 274    CLEAR_EOL = ''       #: Clear to the end of the line.
 275    CLEAR_BOL = ''       #: Clear to the beginning of the line.
 276    CLEAR_EOS = ''       #: Clear to the end of the screen
 277
 278    # Output modes:
 279    BOLD = ''            #: Turn on bold mode
 280    BLINK = ''           #: Turn on blink mode
 281    DIM = ''             #: Turn on half-bright mode
 282    REVERSE = ''         #: Turn on reverse-video mode
 283    NORMAL = ''          #: Turn off all modes
 284
 285    # Cursor display:
 286    HIDE_CURSOR = ''     #: Make the cursor invisible
 287    SHOW_CURSOR = ''     #: Make the cursor visible
 288
 289    # Terminal size:
 290    COLS = None          #: Width of the terminal (None for unknown)
 291    LINES = None         #: Height of the terminal (None for unknown)
 292
 293    # Foreground colors:
 294    BLACK = BLUE = GREEN = CYAN = RED = MAGENTA = YELLOW = WHITE = ''
 295
 296    # Background colors:
 297    BG_BLACK = BG_BLUE = BG_GREEN = BG_CYAN = ''
 298    BG_RED = BG_MAGENTA = BG_YELLOW = BG_WHITE = ''
 299
 300    _STRING_CAPABILITIES = """
 301    BOL=cr UP=cuu1 DOWN=cud1 LEFT=cub1 RIGHT=cuf1
 302    CLEAR_SCREEN=clear CLEAR_EOL=el CLEAR_BOL=el1 CLEAR_EOS=ed BOLD=bold
 303    BLINK=blink DIM=dim REVERSE=rev UNDERLINE=smul NORMAL=sgr0
 304    HIDE_CURSOR=civis SHOW_CURSOR=cnorm""".split()
 305    _COLORS = """BLACK BLUE GREEN CYAN RED MAGENTA YELLOW WHITE""".split()
 306    _ANSICOLORS = "BLACK RED GREEN YELLOW BLUE MAGENTA CYAN WHITE".split()
 307
 308    def __init__(self, term_stream=_sys.stdout):
 309       """
 310       Create a `TerminalController` and initialize its attributes
 311       with appropriate values for the current terminal.
 312       `term_stream` is the stream that will be used for terminal
 313       output; if this stream is not a tty, then the terminal is
 314       assumed to be a dumb terminal (i.e., have no capabilities).
 315       """
 316       # Curses isn't available on all platforms
 317       try: import curses
 318       except: return
 319
 320       # If the stream isn't a tty, then assume it has no capabilities.
 321       if hasattr(_sys.stdout, 'isatty'):
 322            if not term_stream.isatty(): return
 323
 324       # Check the terminal type.  If we fail, then assume that the
 325       # terminal has no capabilities.
 326       try: curses.setupterm()
 327       except: return
 328
 329       # Look up numeric capabilities.
 330       self.COLS = curses.tigetnum('cols')
 331       self.LINES = curses.tigetnum('lines')
 332
 333       # Look up string capabilities.
 334       for capability in self._STRING_CAPABILITIES:
 335          (attrib, cap_name) = capability.split('=')
 336          setattr(self, attrib, self._tigetstr(cap_name) or '')
 337
 338       # Colors
 339       set_fg = self._tigetstr('setf')
 340       if set_fg:
 341          for i,color in enumerate(self._COLORS):
 342             setattr(self, color,
 343                   self._tparm(set_fg, i) or '')
 344       set_fg_ansi = self._tigetstr('setaf')
 345       if set_fg_ansi:
 346          for i,color in enumerate(self._ANSICOLORS):
 347             setattr(self, color,
 348                   self._tparm(set_fg_ansi, i) or '')
 349       set_bg = self._tigetstr('setb')
 350       if set_bg:
 351          for i,color in enumerate(self._COLORS):
 352             setattr(self, 'BG_'+color,
 353                   self._tparm(set_bg, i) or '')
 354       set_bg_ansi = self._tigetstr('setab')
 355       if set_bg_ansi:
 356          for i,color in enumerate(self._ANSICOLORS):
 357             setattr(self, 'BG_'+color,
 358                   self._tparm(set_bg_ansi, i) or '')
 359
 360    def _tigetstr(self, cap_name):
 361       # String capabilities can include "delays" of the form "$<2>".
 362       # For any modern terminal, we should be able to just ignore
 363       # these, so strip them out.
 364       import curses
 365       cap = curses.tigetstr(cap_name)
 366       if cap:
 367          cap = cap.decode()
 368       else:
 369          cap = ''
 370       return _re.sub(r'\$<\d+>[/*]?', '', cap)
 371
 372    def _tparm(self, parm, i):
 373       import curses
 374       res = curses.tparm(parm, i)
 375       if res:
 376          res = res.decode()
 377       return res
 378
 379    def render(self, template):
 380       """
 381       Replace each $-substitutions in the given template string with
 382       the corresponding terminal control string (if it's defined) or
 383       '' (if it's not).
 384       """
 385       return _re.sub(r'\$\$|\${\w+}', self._render_sub, template)
 386
 387    def _render_sub(self, match):
 388       s = match.group()
 389       if s == '$$': return s
 390       else:
 391          return getattr(self, s[2:-1])
 392
 393 def cecho(*args, **kwargs):
 394    """print() replacement which does colored output if supported.
 395
 396    The following sequences will be replaced in the `args`:
 397       * ${BOLD}       : Turn on bold mode
 398       * ${BLINK}      : Turn on blink mode
 399       * ${DIM}        : Turn on half-bright mode
 400       * ${REVERSE}    : Turn on reverse-video mode
 401       * ${NORMAL}     : Turn off all modes
 402       * ${<COLOR>}    : Set text color to <COLOR>
 403       * ${<BG_COLOR>} : Set background color to <BG_COLOR>
 404
 405    The available colors are:
 406       * BLACK
 407       * BLUE
 408       * GREEN
 409       * CYAN
 410       * RED
 411       * MAGENTA
 412       * YELLOW
 413       * WHITE
 414
 415     The valid background colors are:
 416       * BG_BLACK
 417       * BG_BLUE
 418       * BG_GREEN
 419       * BG_CYAN
 420       * BG_RED
 421       * BG_MAGENTA
 422       * BG_YELLOW
 423       * BG_WHITE
 424
 425    """
 426    term = TerminalController()
 427    parts = list(args)
 428    for i, s in enumerate(parts):
 429       if type(s) == str:
 430          parts[i] = term.render(s)
 431    echo(*parts, **kwargs)
 432
 433 def cerror(*args, **kwargs):
 434    """Print a message, preceded with a bold, red '*** Error***' flag"""
 435    cecho("${RED}${BOLD}*** Error ***${NORMAL}", *args, **kwargs)
 436
 437 # ------------------------- vim: set sw=3 sts=3 et: --------------- end-of-file