Prechádzať zdrojové kódy

r.pack/r.unpack: fix pack and unpack VRT raster map with tiled raster maps (#1074)

Tomas Zigo 4 rokov pred
rodič
commit
d96b67d4c1
2 zmenil súbory, kde vykonal 116 pridanie a 23 odobranie
  1. 52 7
      scripts/r.pack/r.pack.py
  2. 64 16
      scripts/r.unpack/r.unpack.py

+ 52 - 7
scripts/r.pack/r.pack.py

@@ -81,16 +81,57 @@ def main():
     olddir = os.getcwd()
 
     # copy elements
-    for element in ['cats', 'cell', 'cellhd', 'colr', 'fcell', 'hist']:
+    info = grass.parse_command('r.info', flags='e', map=infile)
+    vrt_files = {}
+    if info['maptype'] == 'virtual':
+        map_file = grass.find_file(
+            name=infile, element='cell_misc',
+        )
+        if map_file['file']:
+            vrt = os.path.join(map_file['file'], 'vrt')
+            if os.path.exists(vrt):
+                with open(vrt, 'r') as f:
+                    for r in f.readlines():
+                        map, mapset = r.split('@')
+                        map_basedir = os.path.sep.join(
+                            os.path.normpath(
+                                map_file['file'],
+                            ).split(os.path.sep)[:-2],
+                        )
+                        vrt_files[map] = map_basedir
+
+    for element in [
+            'cats', 'cell', 'cellhd', 'cell_misc', 'colr', 'fcell',
+            'hist',
+    ]:
         path = os.path.join(basedir, element, infile)
         if os.path.exists(path):
             grass.debug('copying %s' % path)
-            shutil.copyfile(path,
-                            os.path.join(tmp_dir, element))
-
-    if os.path.exists(os.path.join(basedir, 'cell_misc', infile)):
-        shutil.copytree(os.path.join(basedir, 'cell_misc', infile),
-                        os.path.join(tmp_dir, 'cell_misc'))
+            if os.path.isfile(path):
+                shutil.copyfile(
+                    path, os.path.join(tmp_dir, element),
+                )
+            else:
+                shutil.copytree(
+                    path, os.path.join(tmp_dir, element),
+                )
+
+        # Copy vrt files
+        if vrt_files:
+            for f in vrt_files.keys():
+                f_tmp_dir = os.path.join(tmp, f)
+                if not os.path.exists(f_tmp_dir):
+                    os.mkdir(f_tmp_dir)
+                path = os.path.join(vrt_files[f], element, f)
+                grass.debug("copying vrt file {}".format(path))
+                if os.path.isfile(path):
+                    shutil.copyfile(
+                        path, os.path.join(f_tmp_dir, element),
+                    )
+                else:
+                    shutil.copytree(
+                        path, os.path.join(f_tmp_dir, element),
+                    )
 
     if not os.listdir(tmp_dir):
         grass.fatal(_("No raster map components found"))
@@ -111,6 +152,10 @@ def main():
     else:
         tar = tarfile.TarFile.open(name=outfile_base, mode='w:gz')
     tar.add(infile, recursive=True)
+    if vrt_files:
+        for f in vrt_files.keys():
+            tar.add(f, recursive=True)
+
     tar.close()
     try:
         shutil.move(outfile_base, outfile)

+ 64 - 16
scripts/r.unpack/r.unpack.py

@@ -41,6 +41,7 @@
 #%end
 
 import os
+import re
 import sys
 import shutil
 import tarfile
@@ -73,7 +74,10 @@ def main():
     os.chdir(tmp_dir)
     tar = tarfile.TarFile.open(name=input_base, mode='r')
     try:
-        data_name = tar.getnames()[0]
+        data_names = [
+            tarinfo.name for tarinfo in tar.getmembers()
+            if '/' not in tarinfo.name
+        ]
     except:
         grass.fatal(_("Pack file unreadable"))
 
@@ -92,7 +96,7 @@ def main():
     if options['output']:
         map_name = options['output']
     else:
-        map_name = data_name.split('@')[0]
+        map_name = data_names[0].split('@')[0]
 
     gfile = grass.find_file(name=map_name, element='cell', mapset='.')
     if gfile['file']:
@@ -105,7 +109,7 @@ def main():
     # extract data
     tar.extractall()
     tar.close()
-    os.chdir(data_name)
+    os.chdir(data_names[0])
 
     if os.path.exists('cell'):
         pass
@@ -161,20 +165,64 @@ def main():
                               " use the -o flag to ignore them and use"
                               " current location definition."))
 
+    # Vrt raster map, raster maps postfix e.g. 'map-000-001'
+    r = re.compile(r"[-]\d+(?:\.\d+)?[-]\d+(?:\.\d+)?$")
+
+    maps = []
+    vrt_file = None
     # install in $MAPSET
-    for element in ['cats', 'cell', 'cellhd', 'cell_misc', 'colr', 'fcell', 'hist']:
-        if not os.path.exists(element):
-            continue
-        path = os.path.join(mset_dir, element)
-        if not os.path.exists(path):
-            os.mkdir(path)
-        if element == 'cell_misc':
-            path = os.path.join(mset_dir, element, map_name)
-            if os.path.exists(path):
-                shutil.rmtree(path)
-            shutil.copytree('cell_misc', path)
-        else:
-            shutil.copyfile(element, os.path.join(mset_dir, element, map_name))
+    for index, data in enumerate(data_names):
+        map_postfix = ''
+
+        if index > 0 and data == map_name:
+            map_name = data
+
+        match = re.search(r, data)
+        if match:
+            map_postfix = data[match.span(0)[0]:match.span(0)[1]]
+
+        for element in [
+                'cats', 'cell', 'cellhd', 'cell_misc', 'colr', 'fcell',
+                'hist'
+        ]:
+            src_path = os.path.join(tmp_dir, data, element)
+            if not os.path.exists(src_path):
+                continue
+
+            path = os.path.join(mset_dir, element)
+            if not os.path.exists(path):
+                os.mkdir(path)
+
+            if element == 'cell_misc':
+                if index > 0:
+                    maps.append(
+                        "{map_name}{postfix}@{mapset}".format(
+                            map_name=map_name, postfix=map_postfix,
+                            mapset=gisenv['MAPSET'],
+                        ),
+                    )
+
+                path = os.path.join(
+                    mset_dir, element, map_name + map_postfix,
+                )
+                if index == 0:
+                    vrt_file = os.path.join(path, 'vrt')
+
+                if os.path.exists(path):
+                    shutil.rmtree(path)
+                shutil.copytree(src_path, path)
+            else:
+                shutil.copyfile(
+                    src_path, os.path.join(
+                        mset_dir, element, map_name + map_postfix,
+                    )
+                )
+    # Update vrt file
+    if maps:
+        if vrt_file and os.path.exists(vrt_file):
+            files = '\n'.join(maps)
+            with open(vrt_file, 'w') as f:
+                f.write(files)
 
     grass.message(_('Raster map <{name}> unpacked'.format(name=map_name)))