utils.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465
  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. if sys.version_info.major == 3:
  23. unicode = str
  24. def float_or_dms(s):
  25. """Convert DMS to float.
  26. >>> round(float_or_dms('26:45:30'), 5)
  27. 26.75833
  28. >>> round(float_or_dms('26:0:0.1'), 5)
  29. 26.00003
  30. :param s: DMS value
  31. :return: float value
  32. """
  33. if s[-1] in ['E', 'W', 'N', 'S']:
  34. s = s[:-1]
  35. return sum(float(x) / 60 ** n for (n, x) in enumerate(s.split(':')))
  36. def separator(sep):
  37. """Returns separator from G_OPT_F_SEP appropriately converted
  38. to character.
  39. >>> separator('pipe')
  40. '|'
  41. >>> separator('comma')
  42. ','
  43. If the string does not match any of the separator keywords,
  44. it is returned as is:
  45. >>> separator(', ')
  46. ', '
  47. :param str separator: character or separator keyword
  48. :return: separator character
  49. """
  50. if sep == "pipe":
  51. return "|"
  52. elif sep == "comma":
  53. return ","
  54. elif sep == "space":
  55. return " "
  56. elif sep == "tab" or sep == "\\t":
  57. return "\t"
  58. elif sep == "newline" or sep == "\\n":
  59. return "\n"
  60. return sep
  61. def diff_files(filename_a, filename_b):
  62. """Diffs two text files and returns difference.
  63. :param str filename_a: first file path
  64. :param str filename_b: second file path
  65. :return: list of strings
  66. """
  67. import difflib
  68. differ = difflib.Differ()
  69. fh_a = open(filename_a, 'r')
  70. fh_b = open(filename_b, 'r')
  71. result = list(differ.compare(fh_a.readlines(),
  72. fh_b.readlines()))
  73. return result
  74. def try_remove(path):
  75. """Attempt to remove a file; no exception is generated if the
  76. attempt fails.
  77. :param str path: path to file to remove
  78. """
  79. try:
  80. os.remove(path)
  81. except:
  82. pass
  83. def try_rmdir(path):
  84. """Attempt to remove a directory; no exception is generated if the
  85. attempt fails.
  86. :param str path: path to directory to remove
  87. """
  88. try:
  89. os.rmdir(path)
  90. except:
  91. shutil.rmtree(path, ignore_errors=True)
  92. def basename(path, ext=None):
  93. """Remove leading directory components and an optional extension
  94. from the specified path
  95. :param str path: path
  96. :param str ext: extension
  97. """
  98. name = os.path.basename(path)
  99. if not ext:
  100. return name
  101. fs = name.rsplit('.', 1)
  102. if len(fs) > 1 and fs[1].lower() == ext:
  103. name = fs[0]
  104. return name
  105. class KeyValue(dict):
  106. """A general-purpose key-value store.
  107. KeyValue is a subclass of dict, but also allows entries to be read and
  108. written using attribute syntax. Example:
  109. >>> reg = KeyValue()
  110. >>> reg['north'] = 489
  111. >>> reg.north
  112. 489
  113. >>> reg.south = 205
  114. >>> reg['south']
  115. 205
  116. """
  117. def __getattr__(self, key):
  118. return self[key]
  119. def __setattr__(self, key, value):
  120. self[key] = value
  121. def _get_encoding():
  122. encoding = locale.getdefaultlocale()[1]
  123. if not encoding:
  124. encoding = 'UTF-8'
  125. return encoding
  126. def decode(bytes_, encoding=None):
  127. """Decode bytes with default locale and return (unicode) string
  128. No-op if parameter is not bytes (assumed unicode string).
  129. :param bytes bytes_: the bytes to decode
  130. :param encoding: encoding to be used, default value is None
  131. Example
  132. -------
  133. >>> decode(b'S\xc3\xbcdtirol')
  134. u'Südtirol'
  135. >>> decode(u'Südtirol')
  136. u'Südtirol'
  137. >>> decode(1234)
  138. u'1234'
  139. """
  140. if isinstance(bytes_, unicode):
  141. return bytes_
  142. if isinstance(bytes_, bytes):
  143. if encoding is None:
  144. enc = _get_encoding()
  145. else:
  146. enc = encoding
  147. return bytes_.decode(enc)
  148. # if something else than text
  149. if sys.version_info.major >= 3:
  150. # only text should be used
  151. raise TypeError("can only accept types str and bytes")
  152. else:
  153. # for backwards compatibility
  154. return unicode(bytes_)
  155. def encode(string, encoding=None):
  156. """Encode string with default locale and return bytes with that encoding
  157. No-op if parameter is bytes (assumed already encoded).
  158. This ensures garbage in, garbage out.
  159. :param str string: the string to encode
  160. :param encoding: encoding to be used, default value is None
  161. Example
  162. -------
  163. >>> encode(b'S\xc3\xbcdtirol')
  164. b'S\xc3\xbcdtirol'
  165. >>> decode(u'Südtirol')
  166. b'S\xc3\xbcdtirol'
  167. >>> decode(1234)
  168. b'1234'
  169. """
  170. if isinstance(string, bytes):
  171. return string
  172. # this also tests str in Py3:
  173. if isinstance(string, unicode):
  174. if encoding is None:
  175. enc = _get_encoding()
  176. else:
  177. enc = encoding
  178. return string.encode(enc)
  179. # if something else than text
  180. if sys.version_info.major >= 3:
  181. # only text should be used
  182. raise TypeError("can only accept types str and bytes")
  183. else:
  184. # for backwards compatibility
  185. return bytes(string)
  186. def text_to_string(text, encoding=None):
  187. """Convert text to str. Useful when passing text into environments,
  188. in Python 2 it needs to be bytes on Windows, in Python 3 in needs unicode.
  189. """
  190. if sys.version[0] == '2':
  191. # Python 2
  192. return encode(text, encoding=encoding)
  193. else:
  194. # Python 3
  195. return decode(text, encoding=encoding)
  196. def parse_key_val(s, sep='=', dflt=None, val_type=None, vsep=None):
  197. """Parse a string into a dictionary, where entries are separated
  198. by newlines and the key and value are separated by `sep` (default: `=`)
  199. >>> parse_key_val('min=20\\nmax=50') == {'min': '20', 'max': '50'}
  200. True
  201. >>> parse_key_val('min=20\\nmax=50',
  202. ... val_type=float) == {'min': 20, 'max': 50}
  203. True
  204. :param str s: string to be parsed
  205. :param str sep: key/value separator
  206. :param dflt: default value to be used
  207. :param val_type: value type (None for no cast)
  208. :param vsep: vertical separator (default is Python 'universal newlines' approach)
  209. :return: parsed input (dictionary of keys/values)
  210. """
  211. result = KeyValue()
  212. if not s:
  213. return result
  214. if isinstance(s, bytes):
  215. sep = encode(sep)
  216. vsep = encode(vsep) if vsep else vsep
  217. if vsep:
  218. lines = s.split(vsep)
  219. try:
  220. lines.remove('\n')
  221. except ValueError:
  222. pass
  223. else:
  224. lines = s.splitlines()
  225. for line in lines:
  226. kv = line.split(sep, 1)
  227. k = decode(kv[0].strip())
  228. if len(kv) > 1:
  229. v = decode(kv[1].strip())
  230. else:
  231. v = dflt
  232. if val_type:
  233. result[k] = val_type(v)
  234. else:
  235. result[k] = v
  236. return result
  237. def get_num_suffix(number, max_number):
  238. """Returns formatted number with number of padding zeros
  239. depending on maximum number, used for creating suffix for data series.
  240. Does not include the suffix separator.
  241. :param number: number to be formatted as map suffix
  242. :param max_number: maximum number of the series to get number of digits
  243. >>> get_num_suffix(10, 1000)
  244. '0010'
  245. >>> get_num_suffix(10, 10)
  246. '10'
  247. """
  248. return '{number:0{width}d}'.format(width=len(str(max_number)),
  249. number=number)
  250. def split(s):
  251. """!Platform specific shlex.split"""
  252. if sys.version_info >= (2, 6):
  253. return shlex.split(s, posix = (sys.platform != "win32"))
  254. elif sys.platform == "win32":
  255. return shlex.split(s.replace('\\', r'\\'))
  256. else:
  257. return shlex.split(s)
  258. # source:
  259. # http://stackoverflow.com/questions/4836710/
  260. # does-python-have-a-built-in-function-for-string-natural-sort/4836734#4836734
  261. def natural_sort(l):
  262. """Returns sorted strings using natural sort
  263. """
  264. convert = lambda text: int(text) if text.isdigit() else text.lower()
  265. alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
  266. return sorted(l, key=alphanum_key)
  267. def get_lib_path(modname, libname=None):
  268. """Return the path of the libname contained in the module.
  269. """
  270. from os.path import isdir, join, sep
  271. from os import getenv
  272. if isdir(join(getenv('GISBASE'), 'etc', modname)):
  273. path = join(os.getenv('GISBASE'), 'etc', modname)
  274. elif getenv('GRASS_ADDON_BASE') and libname and \
  275. isdir(join(getenv('GRASS_ADDON_BASE'), 'etc', modname, libname)):
  276. path = join(getenv('GRASS_ADDON_BASE'), 'etc', modname)
  277. elif getenv('GRASS_ADDON_BASE') and \
  278. isdir(join(getenv('GRASS_ADDON_BASE'), 'etc', modname)):
  279. path = join(getenv('GRASS_ADDON_BASE'), 'etc', modname)
  280. elif getenv('GRASS_ADDON_BASE') and \
  281. isdir(join(getenv('GRASS_ADDON_BASE'), modname, modname)):
  282. path = join(os.getenv('GRASS_ADDON_BASE'), modname, modname)
  283. else:
  284. # used by g.extension compilation process
  285. cwd = os.getcwd()
  286. idx = cwd.find(modname)
  287. if idx < 0:
  288. return None
  289. path = '{cwd}{sep}etc{sep}{modname}'.format(cwd=cwd[:idx+len(modname)],
  290. sep=sep,
  291. modname=modname)
  292. if libname:
  293. path += '{pathsep}{cwd}{sep}etc{sep}{modname}{sep}{libname}'.format(
  294. cwd=cwd[:idx+len(modname)],
  295. sep=sep,
  296. modname=modname, libname=libname,
  297. pathsep=os.pathsep
  298. )
  299. return path
  300. def set_path(modulename, dirname=None, path='.'):
  301. """Set sys.path looking in the the local directory GRASS directories.
  302. :param modulename: string with the name of the GRASS module
  303. :param dirname: string with the directory name containing the python
  304. libraries, default None
  305. :param path: string with the path to reach the dirname locally.
  306. Example
  307. --------
  308. "set_path" example working locally with the source code of a module
  309. (r.green) calling the function with all the parameters. Below it is
  310. reported the directory structure on the r.green module.
  311. ::
  312. grass_prompt> pwd
  313. ~/Download/r.green/r.green.hydro/r.green.hydro.financial
  314. grass_prompt> tree ../../../r.green
  315. ../../../r.green
  316. |-- ...
  317. |-- libgreen
  318. | |-- pyfile1.py
  319. | +-- pyfile2.py
  320. +-- r.green.hydro
  321. |-- Makefile
  322. |-- libhydro
  323. | |-- pyfile1.py
  324. | +-- pyfile2.py
  325. |-- r.green.hydro.*
  326. +-- r.green.hydro.financial
  327. |-- Makefile
  328. |-- ...
  329. +-- r.green.hydro.financial.py
  330. 21 directories, 125 files
  331. in the source code the function is called with the following parameters: ::
  332. set_path('r.green', 'libhydro', '..')
  333. set_path('r.green', 'libgreen', os.path.join('..', '..'))
  334. when we are executing the module: r.green.hydro.financial locally from
  335. the command line: ::
  336. grass_prompt> python r.green.hydro.financial.py --ui
  337. In this way we are executing the local code even if the module was already
  338. installed as grass-addons and it is available in GRASS standards path.
  339. The function is cheching if the dirname is provided and if the
  340. directory exists and it is available using the path
  341. provided as third parameter, if yes add the path to sys.path to be
  342. importable, otherwise it will check on GRASS GIS standard paths.
  343. """
  344. import sys
  345. # TODO: why dirname is checked first - the logic should be revised
  346. pathlib = None
  347. if dirname:
  348. pathlib = os.path.join(path, dirname)
  349. if pathlib and os.path.exists(pathlib):
  350. # we are running the script from the script directory, therefore
  351. # we add the path to sys.path to reach the directory (dirname)
  352. sys.path.append(os.path.abspath(path))
  353. else:
  354. # running from GRASS GIS session
  355. path = get_lib_path(modulename, dirname)
  356. if path is None:
  357. pathname = os.path.join(modulename, dirname) if dirname else modulename
  358. raise ImportError("Not able to find the path '%s' directory "
  359. "(current dir '%s')." % (pathname, os.getcwd()))
  360. sys.path.insert(0, path)
  361. def clock():
  362. """
  363. Return time counter to measure performance for chunks of code.
  364. Uses time.clock() for Py < 3.3, time.perf_counter() for Py >= 3.3.
  365. Should be used only as difference between the calls.
  366. """
  367. if sys.version_info > (3,2):
  368. return time.perf_counter()
  369. return time.clock()