瀏覽代碼

pythonlib: safer mechanism for switching dbs in create_location (#804)

* create_location shouldn't switch locations, it uses create_environment instead
* fix create_environment to not use g.tempfile, but Python NamedTemporaryFile
Anna Petrasova 4 年之前
父節點
當前提交
aaf94906c8
共有 1 個文件被更改,包括 17 次插入19 次删除
  1. 17 19
      lib/python/script/core.py

+ 17 - 19
lib/python/script/core.py

@@ -28,9 +28,10 @@ import codecs
 import string
 import random
 import pipes
+from tempfile import NamedTemporaryFile
 import types as python_types
 
-from .utils import KeyValue, parse_key_val, basename, encode, decode
+from .utils import KeyValue, parse_key_val, basename, encode, decode, try_remove
 from grass.exceptions import ScriptError, CalledModuleError
 
 # PY2/PY3 compat
@@ -1588,19 +1589,17 @@ def create_location(dbase, location, epsg=None, proj4=None, filename=None,
     :param bool overwrite: True to overwrite location if exists(WARNING:
                            ALL DATA from existing location ARE DELETED!)
     """
-    gisdbase = None
-    if epsg or proj4 or filename or wkt:
-        # FIXME: changing GISDBASE mid-session is not background-job safe
-        gisdbase = gisenv()['GISDBASE']
-        run_command('g.gisenv', set='GISDBASE=%s' % dbase)
     # create dbase if not exists
     if not os.path.exists(dbase):
-            os.mkdir(dbase)
-
+        os.mkdir(dbase)
+    if epsg or proj4 or filename or wkt:
+        # here the location shouldn't really matter
+        tmp_gisrc, env = create_environment(dbase, gisenv()['LOCATION_NAME'], 'PERMANENT')
     # check if location already exists
     if os.path.exists(os.path.join(dbase, location)):
         if not overwrite:
             warning(_("Location <%s> already exists. Operation canceled.") % location)
+            try_remove(tmp_gisrc)
             return
         else:
             warning(_("Location <%s> already exists and will be overwritten") % location)
@@ -1615,27 +1614,27 @@ def create_location(dbase, location, epsg=None, proj4=None, filename=None,
 
     if epsg:
         ps = pipe_command('g.proj', quiet=True, flags='t', epsg=epsg,
-                          location=location, stderr=PIPE, **kwargs)
+                          location=location, stderr=PIPE, env=env, **kwargs)
     elif proj4:
         ps = pipe_command('g.proj', quiet=True, flags='t', proj4=proj4,
-                          location=location, stderr=PIPE, **kwargs)
+                          location=location, stderr=PIPE, env=env, **kwargs)
     elif filename:
         ps = pipe_command('g.proj', quiet=True, georef=filename,
-                          location=location, stderr=PIPE)
+                          location=location, stderr=PIPE, env=env)
     elif wkt:
         if os.path.isfile(wkt):
             ps = pipe_command('g.proj', quiet=True, wkt=wkt, location=location,
-                              stderr=PIPE)
+                              stderr=PIPE, env=env)
         else:
             ps = pipe_command('g.proj', quiet=True, wkt='-', location=location,
-                              stderr=PIPE, stdin=PIPE)
+                              stderr=PIPE, stdin=PIPE, env=env)
             stdin = encode(wkt)
     else:
         _create_location_xy(dbase, location)
 
     if epsg or proj4 or filename or wkt:
         error = ps.communicate(stdin)[1]
-        run_command('g.gisenv', set='GISDBASE=%s' % gisdbase)
+        try_remove(tmp_gisrc)
 
         if ps.returncode != 0 and error:
             raise ScriptError(repr(error))
@@ -1771,16 +1770,15 @@ def legal_name(s):
 def create_environment(gisdbase, location, mapset):
     """Creates environment to be passed in run_command for example.
     Returns tuple with temporary file path and the environment. The user
-    of this function is responsile for deleting the file."""
-    tmp_gisrc_file = tempfile()
-    with open(tmp_gisrc_file, 'w') as f:
+    of this function is responsible for deleting the file."""
+    with NamedTemporaryFile(mode='w', delete=False) as f:
         f.write('MAPSET: {mapset}\n'.format(mapset=mapset))
         f.write('GISDBASE: {g}\n'.format(g=gisdbase))
         f.write('LOCATION_NAME: {l}\n'.format(l=location))
         f.write('GUI: text\n')
     env = os.environ.copy()
-    env['GISRC'] = tmp_gisrc_file
-    return tmp_gisrc_file, env
+    env['GISRC'] = f.name
+    return f.name, env
 
 
 if __name__ == '__main__':