utils.py 19 KB

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