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