core.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513
  1. """
  2. @package grass.script.core
  3. @brief GRASS Python scripting module
  4. Core functions to be used in Python scripts.
  5. Usage:
  6. @code
  7. from grass.script import core as grass
  8. grass.parser()
  9. ...
  10. @endcode
  11. (C) 2008-2009 by the GRASS Development Team
  12. This program is free software under the GNU General Public
  13. License (>=v2). Read the file COPYING that comes with GRASS
  14. for details.
  15. @author Glynn Clements
  16. @author Martin Landa <landa.martin gmail.com>
  17. """
  18. import os
  19. import sys
  20. import types
  21. import re
  22. import atexit
  23. import subprocess
  24. # subprocess wrapper that uses shell on Windows
  25. class Popen(subprocess.Popen):
  26. def __init__(self, args, bufsize=0, executable=None,
  27. stdin=None, stdout=None, stderr=None,
  28. preexec_fn=None, close_fds=False, shell=None,
  29. cwd=None, env=None, universal_newlines=False,
  30. startupinfo=None, creationflags=0):
  31. if shell == None:
  32. shell = (sys.platform == "win32")
  33. subprocess.Popen.__init__(self, args, bufsize, executable,
  34. stdin, stdout, stderr,
  35. preexec_fn, close_fds, shell,
  36. cwd, env, universal_newlines,
  37. startupinfo, creationflags)
  38. PIPE = subprocess.PIPE
  39. STDOUT = subprocess.STDOUT
  40. def call(*args, **kwargs):
  41. return Popen(*args, **kwargs).wait()
  42. # GRASS-oriented interface to subprocess module
  43. _popen_args = ["bufsize", "executable", "stdin", "stdout", "stderr",
  44. "preexec_fn", "close_fds", "cwd", "env",
  45. "universal_newlines", "startupinfo", "creationflags"]
  46. def _make_val(val):
  47. if isinstance(val, types.StringType):
  48. return val
  49. if isinstance(val, types.ListType):
  50. return ",".join(map(_make_val, val))
  51. if isinstance(val, types.TupleType):
  52. return _make_val(list(val))
  53. return str(val)
  54. def make_command(prog, flags = "", overwrite = False, quiet = False, verbose = False, **options):
  55. """Return a list of strings suitable for use as the args parameter to
  56. Popen() or call(). Example:
  57. >>> grass.make_command("g.message", flags = 'w', message = 'this is a warning')
  58. ['g.message', '-w', 'message=this is a warning']
  59. """
  60. args = [prog]
  61. if overwrite:
  62. args.append("--o")
  63. if quiet:
  64. args.append("--q")
  65. if verbose:
  66. args.append("--v")
  67. if flags:
  68. args.append("-%s" % flags)
  69. for opt, val in options.iteritems():
  70. if val != None:
  71. if opt[0] == '_':
  72. opt = opt[1:]
  73. args.append("%s=%s" % (opt, _make_val(val)))
  74. return args
  75. def start_command(prog, flags = "", overwrite = False, quiet = False, verbose = False, **kwargs):
  76. """Returns a Popen object with the command created by make_command.
  77. Accepts any of the arguments which Popen() accepts apart from "args"
  78. and "shell".
  79. """
  80. options = {}
  81. popts = {}
  82. for opt, val in kwargs.iteritems():
  83. if opt in _popen_args:
  84. popts[opt] = val
  85. else:
  86. options[opt] = val
  87. args = make_command(prog, flags, overwrite, quiet, verbose, **options)
  88. return Popen(args, **popts)
  89. def run_command(*args, **kwargs):
  90. """Passes all arguments to start_command, then waits for the process to
  91. complete, returning its exit code. Similar to subprocess.call(), but
  92. with the make_command() interface.
  93. """
  94. ps = start_command(*args, **kwargs)
  95. return ps.wait()
  96. def pipe_command(*args, **kwargs):
  97. """Passes all arguments to start_command, but also adds
  98. "stdout = PIPE". Returns the Popen object.
  99. """
  100. kwargs['stdout'] = PIPE
  101. return start_command(*args, **kwargs)
  102. def feed_command(*args, **kwargs):
  103. """Passes all arguments to start_command, but also adds
  104. "stdin = PIPE". Returns the Popen object.
  105. """
  106. kwargs['stdin'] = PIPE
  107. return start_command(*args, **kwargs)
  108. def read_command(*args, **kwargs):
  109. """Passes all arguments to pipe_command, then waits for the process to
  110. complete, returning its stdout (i.e. similar to shell `backticks`).
  111. """
  112. ps = pipe_command(*args, **kwargs)
  113. return ps.communicate()[0]
  114. def parse_command(*args, **kwargs):
  115. """Passes all arguments to read_command, then parses the output by
  116. parse_key_val().
  117. Parsing function can be optionally given by <b>parse</b> parameter
  118. including its arguments, e.g.
  119. @code
  120. parse_command(..., parse = (grass.parse_key_val, { 'sep' : ':' }))
  121. @endcode
  122. """
  123. parse = None
  124. if kwargs.has_key('parse'):
  125. if type(kwargs['parse']) is types.TupleType:
  126. parse = kwargs['parse'][0]
  127. parse_args = kwargs['parse'][1]
  128. del kwargs['parse']
  129. if not parse:
  130. parse = parse_key_val # use default fn
  131. parse_args = {}
  132. res = read_command(*args, **kwargs)
  133. return parse(res, **parse_args)
  134. def write_command(*args, **kwargs):
  135. """Passes all arguments to feed_command, with the string specified
  136. by the 'stdin' argument fed to the process' stdin.
  137. """
  138. stdin = kwargs['stdin']
  139. p = feed_command(*args, **kwargs)
  140. p.stdin.write(stdin)
  141. p.stdin.close()
  142. return p.wait()
  143. def exec_command(prog, flags = "", overwrite = False, quiet = False, verbose = False, env = None, **kwargs):
  144. """Interface to os.execvpe(), but with the make_command() interface."""
  145. args = make_command(prog, flags, overwrite, quiet, verbose, **kwargs)
  146. if env == None:
  147. env = os.environ
  148. os.execvpe(prog, args, env)
  149. # interface to g.message
  150. def message(msg, flag = None):
  151. """Display a message using g.message"""
  152. run_command("g.message", flags = flag, message = msg)
  153. def debug(msg, debug = 1):
  154. """Display a debugging message using g.message -d"""
  155. run_command("g.message", flags = 'd', message = msg, debug = debug)
  156. def verbose(msg):
  157. """Display a verbose message using g.message -v"""
  158. message(msg, flag = 'v')
  159. def info(msg):
  160. """Display an informational message using g.message -i"""
  161. message(msg, flag = 'i')
  162. def warning(msg):
  163. """Display a warning message using g.message -w"""
  164. message(msg, flag = 'w')
  165. def error(msg):
  166. """Display an error message using g.message -e"""
  167. message(msg, flag = 'e')
  168. def fatal(msg):
  169. """Display an error message using g.message -e, then abort"""
  170. error(msg)
  171. sys.exit(1)
  172. # interface to g.parser
  173. def _parse_env():
  174. options = {}
  175. flags = {}
  176. for var, val in os.environ.iteritems():
  177. if var.startswith("GIS_OPT_"):
  178. opt = var.replace("GIS_OPT_", "", 1).lower()
  179. options[opt] = val;
  180. if var.startswith("GIS_FLAG_"):
  181. flg = var.replace("GIS_FLAG_", "", 1).lower()
  182. flags[flg] = bool(int(val));
  183. return (options, flags)
  184. def parser():
  185. """Interface to g.parser, intended to be run from the top-level, e.g.:
  186. if __name__ == "__main__":
  187. options, flags = grass.parser()
  188. main()
  189. Thereafter, the global variables "options" and "flags" will be
  190. dictionaries containing option/flag values, keyed by lower-case
  191. option/flag names. The values in "options" are strings, those in
  192. "flags" are Python booleans.
  193. """
  194. if not os.getenv("GISBASE"):
  195. print >> sys.stderr, "You must be in GRASS GIS to run this program."
  196. sys.exit(1)
  197. if len(sys.argv) > 1 and sys.argv[1] == "@ARGS_PARSED@":
  198. return _parse_env()
  199. cmdline = [basename(sys.argv[0])]
  200. cmdline += ['"' + arg + '"' for arg in sys.argv[1:]]
  201. os.environ['CMDLINE'] = ' '.join(cmdline)
  202. argv = sys.argv[:]
  203. name = argv[0]
  204. if not os.path.isabs(name):
  205. if os.sep in name or (os.altsep and os.altsep in name):
  206. argv[0] = os.path.abspath(name)
  207. else:
  208. argv[0] = os.path.join(sys.path[0], name)
  209. if sys.platform == "win32":
  210. os.execvp("g.parser.exe", [name] + argv)
  211. else:
  212. os.execvp("g.parser", [name] + argv)
  213. raise OSError("error executing g.parser")
  214. # interface to g.tempfile
  215. def tempfile():
  216. """Returns the name of a temporary file, created with g.tempfile."""
  217. return read_command("g.tempfile", pid = os.getpid()).strip()
  218. # key-value parsers
  219. def parse_key_val(s, sep = '=', dflt = None, val_type = None, vsep = None):
  220. """Parse a string into a dictionary, where entries are separated
  221. by newlines and the key and value are separated by `sep' (default: `=')
  222. """
  223. result = {}
  224. if not s:
  225. return result
  226. if vsep:
  227. lines = s.split(vsep)
  228. try:
  229. lines.remove('\n')
  230. except ValueError:
  231. pass
  232. else:
  233. lines = s.splitlines()
  234. for line in lines:
  235. kv = line.split(sep, 1)
  236. k = kv[0].strip()
  237. if len(kv) > 1:
  238. v = kv[1]
  239. else:
  240. v = dflt
  241. if val_type:
  242. result[k] = val_type(v)
  243. else:
  244. result[k] = v
  245. return result
  246. # interface to g.gisenv
  247. def gisenv():
  248. """Returns the output from running g.gisenv (with no arguments), as a
  249. dictionary.
  250. """
  251. s = read_command("g.gisenv", flags='n')
  252. return parse_key_val(s)
  253. # interface to g.region
  254. def region():
  255. """Returns the output from running "g.region -g", as a dictionary."""
  256. s = read_command("g.region", flags='g')
  257. return parse_key_val(s, val_type = float)
  258. def use_temp_region():
  259. """Copies the current region to a temporary region with "g.region save=",
  260. then sets WIND_OVERRIDE to refer to that region. Installs an atexit
  261. handler to delete the temporary region upon termination.
  262. """
  263. name = "tmp.%s.%d" % (os.path.basename(sys.argv[0]), os.getpid())
  264. run_command("g.region", save = name)
  265. os.environ['WIND_OVERRIDE'] = name
  266. atexit.register(del_temp_region)
  267. def del_temp_region():
  268. """Unsets WIND_OVERRIDE and removes any region named by it."""
  269. try:
  270. name = os.environ.pop('WIND_OVERRIDE')
  271. run_command("g.remove", quiet = True, region = name)
  272. except:
  273. pass
  274. # interface to g.findfile
  275. def find_file(name, element = 'cell', mapset = None):
  276. """Returns the output from running g.findfile as a dictionary."""
  277. s = read_command("g.findfile", flags='n', element = element, file = name, mapset = mapset)
  278. return parse_key_val(s)
  279. # interface to g.list
  280. def list_grouped(type):
  281. """Returns the output from running g.list, as a dictionary where the keys
  282. are mapset names and the values are lists of maps in that mapset.
  283. """
  284. dashes_re = re.compile("^----+$")
  285. mapset_re = re.compile("<(.*)>")
  286. result = {}
  287. mapset = None
  288. for line in read_command("g.list", type = type).splitlines():
  289. if line == "":
  290. continue
  291. if dashes_re.match(line):
  292. continue
  293. m = mapset_re.search(line)
  294. if m:
  295. mapset = m.group(1)
  296. result[mapset] = []
  297. continue
  298. if mapset:
  299. result[mapset].extend(line.split())
  300. return result
  301. def mlist_grouped(type, mapset = None, pattern = None):
  302. """Returns the output from running g.mlist, as a dictionary where the keys
  303. are mapset names and the values are lists of maps in that mapset.
  304. """
  305. result = {}
  306. mapset_element = None
  307. for line in read_command("g.mlist", flags="m",
  308. type = type, mapset = mapset, pattern = pattern).splitlines():
  309. try:
  310. map, mapset_element = line.split('@')
  311. except ValueError:
  312. print >> sys.stderr, "Invalid element '%s'" % line
  313. continue
  314. if result.has_key(mapset_element):
  315. result[mapset_element].append(map)
  316. else:
  317. result[mapset_element] = [map, ]
  318. return result
  319. def _concat(xs):
  320. result = []
  321. for x in xs:
  322. result.extend(x)
  323. return result
  324. def list_pairs(type):
  325. """Returns the output from running g.list, as a list of (map, mapset)
  326. pairs.
  327. """
  328. return _concat([[(map, mapset) for map in maps]
  329. for mapset, maps in list_grouped(type).iteritems()])
  330. def list_strings(type):
  331. """Returns the output from running g.list, as a list of qualified names."""
  332. return ["%s@%s" % pair for pair in list_pairs(type)]
  333. # color parsing
  334. named_colors = {
  335. "white": (1.00, 1.00, 1.00),
  336. "black": (0.00, 0.00, 0.00),
  337. "red": (1.00, 0.00, 0.00),
  338. "green": (0.00, 1.00, 0.00),
  339. "blue": (0.00, 0.00, 1.00),
  340. "yellow": (1.00, 1.00, 0.00),
  341. "magenta": (1.00, 0.00, 1.00),
  342. "cyan": (0.00, 1.00, 1.00),
  343. "aqua": (0.00, 0.75, 0.75),
  344. "grey": (0.75, 0.75, 0.75),
  345. "gray": (0.75, 0.75, 0.75),
  346. "orange": (1.00, 0.50, 0.00),
  347. "brown": (0.75, 0.50, 0.25),
  348. "purple": (0.50, 0.00, 1.00),
  349. "violet": (0.50, 0.00, 1.00),
  350. "indigo": (0.00, 0.50, 1.00)}
  351. def parse_color(val, dflt = None):
  352. """Parses the string "val" as a GRASS colour, which can be either one of
  353. the named colours or an R:G:B tuple e.g. 255:255:255. Returns an
  354. (r,g,b) triple whose components are floating point values between 0
  355. and 1.
  356. """
  357. if val in named_colors:
  358. return named_colors[val]
  359. vals = val.split(':')
  360. if len(vals) == 3:
  361. return tuple(float(v) / 255 for v in vals)
  362. return dflt
  363. # check GRASS_OVERWRITE
  364. def overwrite():
  365. """Return True if existing files may be overwritten"""
  366. owstr = 'GRASS_OVERWRITE'
  367. return owstr in os.environ and os.environ[owstr] != '0'
  368. # check GRASS_VERBOSE
  369. def verbosity():
  370. """Return the verbosity level selected by GRASS_VERBOSE"""
  371. vbstr = os.getenv('GRASS_VERBOSE')
  372. if vbstr:
  373. return int(vbstr)
  374. else:
  375. return 0
  376. ## various utilities, not specific to GRASS
  377. # basename inc. extension stripping
  378. def basename(path, ext = None):
  379. """Remove leading directory components and an optional extension
  380. from the specified path
  381. """
  382. name = os.path.basename(path)
  383. if not ext:
  384. return name
  385. fs = name.rsplit('.', 1)
  386. if len(fs) > 1 and fs[1].lower() == ext:
  387. name = fs[0]
  388. return name
  389. # find a program (replacement for "which")
  390. def find_program(pgm, args = []):
  391. """Attempt to run a program, with optional arguments. Return False
  392. if the attempt failed due to a missing executable, True otherwise
  393. """
  394. nuldev = file(os.devnull, 'w+')
  395. try:
  396. call([pgm] + args, stdin = nuldev, stdout = nuldev, stderr = nuldev)
  397. found = True
  398. except:
  399. found = False
  400. nuldev.close()
  401. return found
  402. # try to remove a file, without complaints
  403. def try_remove(path):
  404. """Attempt to remove a file; no exception is generated if the
  405. attempt fails.
  406. """
  407. try:
  408. os.remove(path)
  409. except:
  410. pass
  411. # try to remove a directory, without complaints
  412. def try_rmdir(path):
  413. """Attempt to remove a directory; no exception is generated if the
  414. attempt fails.
  415. """
  416. try:
  417. os.rmdir(path)
  418. except:
  419. pass
  420. def float_or_dms(s):
  421. return sum(float(x) / 60 ** n for (n, x) in enumerate(s.split(':')))