utils.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630
  1. """
  2. Useful functions to be used in Python scripts.
  3. Usage:
  4. ::
  5. from grass.script import utils as gutils
  6. (C) 2014-2016 by the GRASS Development Team
  7. This program is free software under the GNU General Public
  8. License (>=v2). Read the file COPYING that comes with GRASS
  9. for details.
  10. .. sectionauthor:: Glynn Clements
  11. .. sectionauthor:: Martin Landa <landa.martin gmail.com>
  12. .. sectionauthor:: Anna Petrasova <kratochanna gmail.com>
  13. """
  14. import os
  15. import sys
  16. import shutil
  17. import locale
  18. import shlex
  19. import re
  20. import time
  21. import platform
  22. import uuid
  23. import random
  24. import string
  25. if sys.version_info.major >= 3:
  26. unicode = str
  27. def float_or_dms(s):
  28. """Convert DMS to float.
  29. >>> round(float_or_dms('26:45:30'), 5)
  30. 26.75833
  31. >>> round(float_or_dms('26:0:0.1'), 5)
  32. 26.00003
  33. :param s: DMS value
  34. :return: float value
  35. """
  36. if s[-1] in ["E", "W", "N", "S"]:
  37. s = s[:-1]
  38. return sum(float(x) / 60**n for (n, x) in enumerate(s.split(":")))
  39. def separator(sep):
  40. """Returns separator from G_OPT_F_SEP appropriately converted
  41. to character.
  42. >>> separator('pipe')
  43. '|'
  44. >>> separator('comma')
  45. ','
  46. If the string does not match any of the separator keywords,
  47. it is returned as is:
  48. >>> separator(', ')
  49. ', '
  50. :param str separator: character or separator keyword
  51. :return: separator character
  52. """
  53. if sep == "pipe":
  54. return "|"
  55. elif sep == "comma":
  56. return ","
  57. elif sep == "space":
  58. return " "
  59. elif sep == "tab" or sep == "\\t":
  60. return "\t"
  61. elif sep == "newline" or sep == "\\n":
  62. return "\n"
  63. return sep
  64. def diff_files(filename_a, filename_b):
  65. """Diffs two text files and returns difference.
  66. :param str filename_a: first file path
  67. :param str filename_b: second file path
  68. :return: list of strings
  69. """
  70. import difflib
  71. differ = difflib.Differ()
  72. fh_a = open(filename_a, "r")
  73. fh_b = open(filename_b, "r")
  74. result = list(differ.compare(fh_a.readlines(), fh_b.readlines()))
  75. return result
  76. def try_remove(path):
  77. """Attempt to remove a file; no exception is generated if the
  78. attempt fails.
  79. :param str path: path to file to remove
  80. """
  81. try:
  82. os.remove(path)
  83. except Exception:
  84. pass
  85. def try_rmdir(path):
  86. """Attempt to remove a directory; no exception is generated if the
  87. attempt fails.
  88. :param str path: path to directory to remove
  89. """
  90. try:
  91. os.rmdir(path)
  92. except Exception:
  93. shutil.rmtree(path, ignore_errors=True)
  94. def basename(path, ext=None):
  95. """Remove leading directory components and an optional extension
  96. from the specified path
  97. :param str path: path
  98. :param str ext: extension
  99. """
  100. name = os.path.basename(path)
  101. if not ext:
  102. return name
  103. fs = name.rsplit(".", 1)
  104. if len(fs) > 1 and fs[1].lower() == ext:
  105. name = fs[0]
  106. return name
  107. class KeyValue(dict):
  108. """A general-purpose key-value store.
  109. KeyValue is a subclass of dict, but also allows entries to be read and
  110. written using attribute syntax. Example:
  111. >>> reg = KeyValue()
  112. >>> reg['north'] = 489
  113. >>> reg.north
  114. 489
  115. >>> reg.south = 205
  116. >>> reg['south']
  117. 205
  118. """
  119. def __getattr__(self, key):
  120. return self[key]
  121. def __setattr__(self, key, value):
  122. self[key] = value
  123. def _get_encoding():
  124. encoding = locale.getdefaultlocale()[1]
  125. if not encoding:
  126. encoding = "UTF-8"
  127. return encoding
  128. def decode(bytes_, encoding=None):
  129. """Decode bytes with default locale and return (unicode) string
  130. No-op if parameter is not bytes (assumed unicode string).
  131. :param bytes bytes_: the bytes to decode
  132. :param encoding: encoding to be used, default value is None
  133. Example
  134. -------
  135. >>> decode(b'S\xc3\xbcdtirol')
  136. u'Südtirol'
  137. >>> decode(u'Südtirol')
  138. u'Südtirol'
  139. >>> decode(1234)
  140. u'1234'
  141. """
  142. if isinstance(bytes_, unicode):
  143. return bytes_
  144. if isinstance(bytes_, bytes):
  145. if encoding is None:
  146. enc = _get_encoding()
  147. else:
  148. enc = encoding
  149. return bytes_.decode(enc)
  150. # if something else than text
  151. if sys.version_info.major >= 3:
  152. # only text should be used
  153. raise TypeError("can only accept types str and bytes")
  154. else:
  155. # for backwards compatibility
  156. return unicode(bytes_)
  157. def encode(string, encoding=None):
  158. """Encode string with default locale and return bytes with that encoding
  159. No-op if parameter is bytes (assumed already encoded).
  160. This ensures garbage in, garbage out.
  161. :param str string: the string to encode
  162. :param encoding: encoding to be used, default value is None
  163. Example
  164. -------
  165. >>> encode(b'S\xc3\xbcdtirol')
  166. b'S\xc3\xbcdtirol'
  167. >>> decode(u'Südtirol')
  168. b'S\xc3\xbcdtirol'
  169. >>> decode(1234)
  170. b'1234'
  171. """
  172. if isinstance(string, bytes):
  173. return string
  174. # this also tests str in Py3:
  175. if isinstance(string, unicode):
  176. if encoding is None:
  177. enc = _get_encoding()
  178. else:
  179. enc = encoding
  180. return string.encode(enc)
  181. # if something else than text
  182. if sys.version_info.major >= 3:
  183. # only text should be used
  184. raise TypeError("can only accept types str and bytes")
  185. else:
  186. # for backwards compatibility
  187. return bytes(string)
  188. def text_to_string(text, encoding=None):
  189. """Convert text to str. Useful when passing text into environments,
  190. in Python 2 it needs to be bytes on Windows, in Python 3 in needs unicode.
  191. """
  192. if sys.version[0] == "2":
  193. # Python 2
  194. return encode(text, encoding=encoding)
  195. else:
  196. # Python 3
  197. return decode(text, encoding=encoding)
  198. def parse_key_val(s, sep="=", dflt=None, val_type=None, vsep=None):
  199. """Parse a string into a dictionary, where entries are separated
  200. by newlines and the key and value are separated by `sep` (default: `=`)
  201. >>> parse_key_val('min=20\\nmax=50') == {'min': '20', 'max': '50'}
  202. True
  203. >>> parse_key_val('min=20\\nmax=50',
  204. ... val_type=float) == {'min': 20, 'max': 50}
  205. True
  206. :param str s: string to be parsed
  207. :param str sep: key/value separator
  208. :param dflt: default value to be used
  209. :param val_type: value type (None for no cast)
  210. :param vsep: vertical separator (default is Python 'universal newlines' approach)
  211. :return: parsed input (dictionary of keys/values)
  212. """
  213. result = KeyValue()
  214. if not s:
  215. return result
  216. if isinstance(s, bytes):
  217. sep = encode(sep)
  218. vsep = encode(vsep) if vsep else vsep
  219. if vsep:
  220. lines = s.split(vsep)
  221. try:
  222. lines.remove("\n")
  223. except ValueError:
  224. pass
  225. else:
  226. lines = s.splitlines()
  227. for line in lines:
  228. kv = line.split(sep, 1)
  229. k = decode(kv[0].strip())
  230. if len(kv) > 1:
  231. v = decode(kv[1].strip())
  232. else:
  233. v = dflt
  234. if val_type:
  235. result[k] = val_type(v)
  236. else:
  237. result[k] = v
  238. return result
  239. def get_num_suffix(number, max_number):
  240. """Returns formatted number with number of padding zeros
  241. depending on maximum number, used for creating suffix for data series.
  242. Does not include the suffix separator.
  243. :param number: number to be formatted as map suffix
  244. :param max_number: maximum number of the series to get number of digits
  245. >>> get_num_suffix(10, 1000)
  246. '0010'
  247. >>> get_num_suffix(10, 10)
  248. '10'
  249. """
  250. return "{number:0{width}d}".format(width=len(str(max_number)), number=number)
  251. def split(s):
  252. """!Platform specific shlex.split"""
  253. if sys.version_info >= (2, 6):
  254. return shlex.split(s, posix=(sys.platform != "win32"))
  255. elif sys.platform == "win32":
  256. return shlex.split(s.replace("\\", r"\\"))
  257. else:
  258. return shlex.split(s)
  259. # source:
  260. # http://stackoverflow.com/questions/4836710/
  261. # does-python-have-a-built-in-function-for-string-natural-sort/4836734#4836734
  262. def natural_sort(items):
  263. """Returns sorted list using natural sort
  264. (deprecated, use naturally_sorted)
  265. """
  266. return naturally_sorted(items)
  267. def naturally_sorted(items, key=None):
  268. """Returns sorted list using natural sort"""
  269. copy_items = items[:]
  270. naturally_sort(copy_items, key)
  271. return copy_items
  272. def naturally_sort(items, key=None):
  273. """Sorts lists using natural sort"""
  274. def convert(text):
  275. return int(text) if text.isdigit() else text.lower()
  276. def alphanum_key(actual_key):
  277. if key:
  278. sort_key = key(actual_key)
  279. else:
  280. sort_key = actual_key
  281. return [convert(c) for c in re.split("([0-9]+)", sort_key)]
  282. items.sort(key=alphanum_key)
  283. def get_lib_path(modname, libname=None):
  284. """Return the path of the libname contained in the module."""
  285. from os.path import isdir, join, sep
  286. from os import getenv
  287. if isdir(join(getenv("GISBASE"), "etc", modname)):
  288. path = join(os.getenv("GISBASE"), "etc", modname)
  289. elif (
  290. getenv("GRASS_ADDON_BASE")
  291. and libname
  292. and isdir(join(getenv("GRASS_ADDON_BASE"), "etc", modname, libname))
  293. ):
  294. path = join(getenv("GRASS_ADDON_BASE"), "etc", modname)
  295. elif getenv("GRASS_ADDON_BASE") and isdir(
  296. join(getenv("GRASS_ADDON_BASE"), "etc", modname)
  297. ):
  298. path = join(getenv("GRASS_ADDON_BASE"), "etc", modname)
  299. elif getenv("GRASS_ADDON_BASE") and isdir(
  300. join(getenv("GRASS_ADDON_BASE"), modname, modname)
  301. ):
  302. path = join(os.getenv("GRASS_ADDON_BASE"), modname, modname)
  303. else:
  304. # used by g.extension compilation process
  305. cwd = os.getcwd()
  306. idx = cwd.find(modname)
  307. if idx < 0:
  308. return None
  309. path = "{cwd}{sep}etc{sep}{modname}".format(
  310. cwd=cwd[: idx + len(modname)], sep=sep, modname=modname
  311. )
  312. if libname:
  313. path += "{pathsep}{cwd}{sep}etc{sep}{modname}{sep}{libname}".format(
  314. cwd=cwd[: idx + len(modname)],
  315. sep=sep,
  316. modname=modname,
  317. libname=libname,
  318. pathsep=os.pathsep,
  319. )
  320. return path
  321. def set_path(modulename, dirname=None, path="."):
  322. """Set sys.path looking in the the local directory GRASS directories.
  323. :param modulename: string with the name of the GRASS module
  324. :param dirname: string with the directory name containing the python
  325. libraries, default None
  326. :param path: string with the path to reach the dirname locally.
  327. Example
  328. --------
  329. "set_path" example working locally with the source code of a module
  330. (r.green) calling the function with all the parameters. Below it is
  331. reported the directory structure on the r.green module.
  332. ::
  333. grass_prompt> pwd
  334. ~/Download/r.green/r.green.hydro/r.green.hydro.financial
  335. grass_prompt> tree ../../../r.green
  336. ../../../r.green
  337. |-- ...
  338. |-- libgreen
  339. | |-- pyfile1.py
  340. | +-- pyfile2.py
  341. +-- r.green.hydro
  342. |-- Makefile
  343. |-- libhydro
  344. | |-- pyfile1.py
  345. | +-- pyfile2.py
  346. |-- r.green.hydro.*
  347. +-- r.green.hydro.financial
  348. |-- Makefile
  349. |-- ...
  350. +-- r.green.hydro.financial.py
  351. 21 directories, 125 files
  352. in the source code the function is called with the following parameters: ::
  353. set_path('r.green', 'libhydro', '..')
  354. set_path('r.green', 'libgreen', os.path.join('..', '..'))
  355. when we are executing the module: r.green.hydro.financial locally from
  356. the command line: ::
  357. grass_prompt> python r.green.hydro.financial.py --ui
  358. In this way we are executing the local code even if the module was already
  359. installed as grass-addons and it is available in GRASS standards path.
  360. The function is cheching if the dirname is provided and if the
  361. directory exists and it is available using the path
  362. provided as third parameter, if yes add the path to sys.path to be
  363. importable, otherwise it will check on GRASS GIS standard paths.
  364. """
  365. import sys
  366. # TODO: why dirname is checked first - the logic should be revised
  367. pathlib = None
  368. if dirname:
  369. pathlib = os.path.join(path, dirname)
  370. if pathlib and os.path.exists(pathlib):
  371. # we are running the script from the script directory, therefore
  372. # we add the path to sys.path to reach the directory (dirname)
  373. sys.path.append(os.path.abspath(path))
  374. else:
  375. # running from GRASS GIS session
  376. path = get_lib_path(modulename, dirname)
  377. if path is None:
  378. pathname = os.path.join(modulename, dirname) if dirname else modulename
  379. raise ImportError(
  380. "Not able to find the path '%s' directory "
  381. "(current dir '%s')." % (pathname, os.getcwd())
  382. )
  383. sys.path.insert(0, path)
  384. def clock():
  385. """
  386. Return time counter to measure performance for chunks of code.
  387. Uses time.clock() for Py < 3.3, time.perf_counter() for Py >= 3.3.
  388. Should be used only as difference between the calls.
  389. """
  390. if sys.version_info > (3, 2):
  391. return time.perf_counter()
  392. return time.clock()
  393. def legalize_vector_name(name, fallback_prefix="x"):
  394. """Make *name* usable for vectors, tables, and columns
  395. The returned string is a name usable for vectors, tables, and columns,
  396. i.e., it is a vector legal name which is a string containing only
  397. lowercase and uppercase ASCII letters, digits, and underscores.
  398. Invalid characters are replaced by underscores.
  399. If the name starts with an invalid character, the name is prefixed with
  400. *fallback_prefix*. This increases the length of the resulting name by the
  401. length of the prefix.
  402. The *fallback_prefix* can be empty which is useful when the *name* is later
  403. used as a suffix for some other valid name.
  404. ValueError is raised when provided *name* is empty or *fallback_prefix*
  405. does not start with a valid character.
  406. """
  407. # The implementation is based on Vect_legal_filename().
  408. if not name:
  409. raise ValueError("name cannot be empty")
  410. if fallback_prefix and re.match("[^A-Za-z]", fallback_prefix[0]):
  411. raise ValueError("fallback_prefix must start with an ASCII letter")
  412. if fallback_prefix and re.match("[^A-Za-z]", name[0], flags=re.ASCII):
  413. # We prefix here rather than just replace, because in cases of unique
  414. # identifiers, e.g., columns or node names, replacing the first
  415. # character by the same replacement character increases chances of
  416. # conflict (e.g. column names 10, 20, 30).
  417. name = "{fallback_prefix}{name}".format(**locals())
  418. name = re.sub("[^A-Za-z0-9_]", "_", name, flags=re.ASCII)
  419. keywords = ["and", "or", "not"]
  420. if name in keywords:
  421. name = "{name}_".format(**locals())
  422. return name
  423. def append_node_pid(name):
  424. """Add node name and PID to a name (string)
  425. For the result to be unique, the name needs to be unique within a process.
  426. Given that, the result will be unique enough for use in temporary maps
  427. and other elements on single machine or an HPC cluster.
  428. The returned string is a name usable for vectors, tables, and columns
  429. (vector legal name) as long as provided argument *name* is.
  430. >>> append_node_pid("tmp_raster_1")
  431. ..note::
  432. Before you use this function for creating temporary files (i.e., normal
  433. files on disk, not maps and other mapset elements), see functions
  434. designed for it in the GRASS GIS or standard Python library. These
  435. take care of collisions already on different levels.
  436. """
  437. # We are using this node as a suffix, so we don't need to make sure it
  438. # is prefixed with additional character(s) since that's exactly what
  439. # happens in this function.
  440. # Note that this may still cause collisions when nodes are named in a way
  441. # that they collapse into the same name after the replacements are done,
  442. # but we consider that unlikely given that
  443. # nodes will be likely already named as something close to what we need.
  444. node = legalize_vector_name(platform.node(), fallback_prefix="")
  445. pid = os.getpid()
  446. return "{name}_{node}_{pid}".format(**locals())
  447. def append_uuid(name):
  448. """Add UUID4 to a name (string)
  449. To generate a name of an temporary mapset element which is unique in a
  450. system, use :func:`append_node_pid()` in a combination with a name unique
  451. within your process.
  452. To avoid collisions, never shorten the name obtained from this function.
  453. A shortened UUID does not have the collision guarantees the full UUID has.
  454. For a random name of a given shorter size, see :func:`append_random()`.
  455. >>> append_uuid("tmp")
  456. ..note::
  457. See the note about creating temporary files in the
  458. :func:`append_node_pid()` description.
  459. """
  460. suffix = uuid.uuid4().hex
  461. return "{name}_{suffix}".format(**locals())
  462. def append_random(name, suffix_length=None, total_length=None):
  463. """Add a random part to of a specified length to a name (string)
  464. >>> append_random("tmp", 8)
  465. >>> append_random("tmp", total_length=16)
  466. ..note::
  467. Note that this will be influeced by the random seed set for the Python
  468. random package.
  469. ..note::
  470. See the note about creating temporary files in the
  471. :func:`append_node_pid()` description.
  472. """
  473. if suffix_length and total_length:
  474. raise ValueError(
  475. "Either suffix_length or total_length can be provided, not both"
  476. )
  477. if not suffix_length and not total_length:
  478. raise ValueError("suffix_length or total_length has to be provided")
  479. if total_length:
  480. # remove len of name and one underscore
  481. name_length = len(name)
  482. suffix_length = total_length - name_length - 1
  483. if suffix_length <= 0:
  484. raise ValueError(
  485. "No characters left for the suffix:"
  486. " total_length <{total_length}> is too small"
  487. " or name <{name}> ({name_length}) is too long".format(**locals())
  488. )
  489. # We don't do lower and upper case because that could cause conflicts in
  490. # contexts which are case-insensitive.
  491. # We use lowercase because that's what is in UUID4 hex string.
  492. allowed_chars = string.ascii_lowercase + string.digits
  493. # The following can be shorter with random.choices from Python 3.6.
  494. suffix = "".join(random.choice(allowed_chars) for _ in range(suffix_length))
  495. return "{name}_{suffix}".format(**locals())