Merge pull request #347 from johntruckenbrodt/feature/archive.select_return-value

[Archive.select] new argument 'return_value'

Author: johntruckenbrodt

environment.yml

@@ -16,6 +16,7 @@ dependencies:
   - python>=3.8
   - pyyaml
   - requests
+  - shapely
   - spatialist>=0.15.2
   - sqlalchemy>=1.4,<2.0
   - sqlalchemy-utils>=0.37

pyroSAR/drivers.py

@@ -2796,7 +2796,7 @@ def cleanup(self):
 
     @staticmethod
     def encode(string, encoding='utf-8'):
-        if not isinstance(string, str):
+        if not isinstance(string, str) and hasattr(string, 'encode'):
             return string.encode(encoding)
         else:
             return string
@@ -3042,7 +3042,7 @@ def move(self, scenelist, directory, pbar=False):
             log.info('The following scenes already exist at the target location:\n{}'.format('\n'.join(double)))
 
     def select(self, vectorobject=None, mindate=None, maxdate=None, date_strict=True,
-               processdir=None, recursive=False, polarizations=None, **args):
+               processdir=None, recursive=False, polarizations=None, return_value="scene", **args):
         """
         select scenes from the database
 
@@ -3057,7 +3057,7 @@ def select(self, vectorobject=None, mindate=None, maxdate=None, date_strict=True
         date_strict: bool
             treat dates as strict limits or also allow flexible limits to incorporate scenes
             whose acquisition period overlaps with the defined limit?
-            
+
             - strict: start >= mindate & stop <= maxdate
             - not strict: stop >= mindate & start <= maxdate
         processdir: str or None
@@ -3067,15 +3067,55 @@ def select(self, vectorobject=None, mindate=None, maxdate=None, date_strict=True
             (only if `processdir` is not None) should also the subdirectories of the `processdir` be scanned?
         polarizations: list[str] or None
             a list of polarization strings, e.g. ['HH', 'VV']
+        return_value: str or List[str]
+            the query return value(s). Options:
+            
+            - geometry_wkb: the scene's footprint geometry formatted as WKB
+            - geometry_wkt: the scene's footprint geometry formatted as WKT
+            - mindate: the acquisition start datetime in UTC formatted as YYYYmmddTHHMMSS
+            - maxdate: the acquisition end datetime in UTC formatted as YYYYmmddTHHMMSS
+            - all further database column names (see :meth:`~Archive.get_colnames()`)
         **args:
             any further arguments (columns), which are registered in the database. See :meth:`~Archive.get_colnames()`
 
         Returns
         -------
-        list[str]
-            the file names pointing to the selected scenes
-
+        List[str] or List[tuple[str]]
+            If a single return_value is specified: list of values for that attribute
+            If multiple return_values are specified: list of tuples containing the requested attributes
         """
+        # Convert return_value to list if it's a string
+        if isinstance(return_value, str):
+            return_values = [return_value]
+        else:
+            return_values = return_value
+        
+        return_values_sql = []
+        for val in return_values:
+            if val == 'mindate':
+                return_values_sql.append('start')
+            elif val == 'maxdate':
+                return_values_sql.append('stop')
+            elif val == 'geometry_wkt':
+                prefix = 'ST_' if self.driver == 'postgresql' else ''
+                return_values_sql.append(f'{prefix}AsText(geometry) as geometry_wkt')
+            elif val == 'geometry_wkb':
+                prefix = 'ST_' if self.driver == 'postgresql' else ''
+                return_values_sql.append(f'{prefix}AsBinary(geometry) as geometry_wkb')
+            else:
+                return_values_sql.append(val)
+        
+        # Validate that all requested return values exist in the database
+        valid_columns = self.get_colnames()
+        extra = ['mindate', 'maxdate', 'geometry_wkt', 'geometry_wkb']
+        normal_returns = [x for x in return_values if x not in extra]
+        invalid_returns = [x for x in normal_returns if x not in valid_columns]
+        if invalid_returns:
+            invalid_str = ', '.join(invalid_returns)
+            msg = (f"The following options are not supported as "
+                   f"return values: {invalid_str}")
+            raise ValueError(msg)
+        
         arg_valid = [x for x in args.keys() if x in self.get_colnames()]
         arg_invalid = [x for x in args.keys() if x not in self.get_colnames()]
         if len(arg_invalid) > 0:
@@ -3091,6 +3131,7 @@ def select(self, vectorobject=None, mindate=None, maxdate=None, date_strict=True
                     arg_format.append("""{0}='{1}'""".format(key, args[key]))
                 elif isinstance(args[key], (tuple, list)):
                     arg_format.append("""{0} IN ('{1}')""".format(key, "', '".join(map(str, args[key]))))
+        
         if mindate:
             if isinstance(mindate, datetime):
                 mindate = mindate.strftime('%Y%m%dT%H%M%S')
@@ -3102,6 +3143,7 @@ def select(self, vectorobject=None, mindate=None, maxdate=None, date_strict=True
                 vals.append(mindate)
             else:
                 log.info('WARNING: argument mindate is ignored, must be in format YYYYmmddTHHMMSS')
+        
         if maxdate:
             if isinstance(maxdate, datetime):
                 maxdate = maxdate.strftime('%Y%m%dT%H%M%S')
@@ -3140,7 +3182,10 @@ def select(self, vectorobject=None, mindate=None, maxdate=None, date_strict=True
             subquery = ' WHERE {}'.format(' AND '.join(arg_format))
         else:
             subquery = ''
-        query = '''SELECT scene, outname_base FROM data{}'''.format(subquery)
+        
+        # Modify the query to select the requested return values
+        query = 'SELECT {}, outname_base FROM data{}'.format(', '.join(return_values_sql), subquery)
+        
         # the query gets assembled stepwise here
         for val in vals:
             query = query.replace('?', """'{0}'""", 1).format(val)
@@ -3151,12 +3196,18 @@ def select(self, vectorobject=None, mindate=None, maxdate=None, date_strict=True
 
         if processdir and os.path.isdir(processdir):
             scenes = [x for x in query_rs
-                      if len(finder(processdir, [x[1]], regex=True, recursive=recursive)) == 0]
+                      if len(finder(processdir, [x[-1]], regex=True, recursive=recursive)) == 0]
         else:
             scenes = query_rs
+        
         ret = []
         for x in scenes:
-            ret.append(self.encode(x[0]))
+            # If only one return value was requested, append just that value
+            if len(return_values) == 1:
+                ret.append(self.encode(x[0]))
+            else:
+                # If multiple return values were requested, append a tuple of all values
+                ret.append(tuple(self.encode(val) for val in x[:-1]))  # Exclude outname_base
 
         return ret
 

requirements.txt

@@ -7,6 +7,7 @@ progressbar2
 psycopg2
 pyyaml
 requests
+shapely
 spatialist>=0.15.2
 sqlalchemy>=1.4,<2.0
 sqlalchemy-utils>=0.37

tests/test_drivers.py

@@ -8,6 +8,7 @@
 from spatialist import Vector
 from sqlalchemy import Table, MetaData, Column, Integer, String
 from geoalchemy2 import Geometry
+from shapely import wkt
 
 metadata = MetaData()
 
@@ -177,6 +178,20 @@ def test_archive(tmpdir, testdata):
     assert len(out) == 1
     assert isinstance(out[0], str)
 
+    out = db.select(vv=1, return_value=['mindate', 'geometry_wkt', 'geometry_wkb'])
+    assert len(out) == 1
+    assert isinstance(out[0], tuple)
+    assert out[0][0] == '20150222T170750'
+    geom = wkt.loads('POLYGON(('
+                     '8.505644 50.295261, 12.0268 50.688881, '
+                     '11.653832 52.183979, 8.017178 51.788181, '
+                     '8.505644 50.295261))')
+    assert wkt.loads(out[0][1]) == geom
+    assert out[0][2] == geom.wkb
+    
+    with pytest.raises(ValueError):
+        out = db.select(vv=1, return_value=['foobar'])
+    
     db.insert(testdata['s1_3'])
     db.insert(testdata['s1_4'])
     db.drop_element(testdata['s1_3'])
@@ -202,9 +217,11 @@ def test_archive2(tmpdir, testdata):
     assert not os.path.isfile(dbfile)
     assert Vector(shp).nfeatures == 1
 
-    with pytest.raises(OSError):
-        with pyroSAR.Archive(dbfile) as db:
+    with pyroSAR.Archive(dbfile) as db:
+        with pytest.raises(OSError):
             db.import_outdated(testdata['archive_old_csv'])
+        with pytest.raises(RuntimeError):
+            db.import_outdated('foobar')
 
     # the archive_old_bbox database contains a relative file name for the scene
     # so that it can be reimported into the new database. The working directory
@@ -249,6 +266,15 @@ def test_archive_postgres(tmpdir, testdata):
     out = db.select(vv=1, acquisition_mode=('IW', 'EW'))
     assert len(out) == 1
     assert isinstance(out[0], str)
+    
+    out = db.select(vv=1, return_value=['scene', 'start'])
+    assert len(out) == 1
+    assert isinstance(out[0], tuple)
+    assert out[0][1] == '20150222T170750'
+    
+    with pytest.raises(ValueError):
+        out = db.select(vv=1, return_value=['foobar'])
+    
     db.add_tables(mytable)
     assert 'mytable' in db.get_tablenames()
     with pytest.raises(TypeError):