[API] method to retrieve corresponding GDAL driver from a file

extension
qgis · Aug 1, 2017 · 0476d4e · 0476d4e
1 parent fef93fc
commit 0476d4e
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 0 deletions.
diff --git a/python/core/raster/qgsrasterfilewriter.sip b/python/core/raster/qgsrasterfilewriter.sip
@@ -156,6 +156,15 @@ class QgsRasterFileWriter
  :rtype: list of str
 %End
 
+    static QString driverForExtension( const QString &extension );
+%Docstring
+ Returns the GDAL driver name for a specified file ``extension``. E.g. the
+ driver name for the ".tif" extension is "GTiff".
+ If no suitable drivers are found then an empty string is returned.
+.. versionadded:: 3.0
+ :rtype: str
+%End
+
 };
 
 /************************************************************************

diff --git a/src/core/raster/qgsrasterfilewriter.cpp b/src/core/raster/qgsrasterfilewriter.cpp
@@ -29,6 +29,9 @@
 #include <QTextStream>
 #include <QMessageBox>
 
+#include <gdal.h>
+#include <cpl_string.h>
+
 QgsRasterDataProvider *QgsRasterFileWriter::createOneBandRaster( Qgis::DataType dataType, int width, int height, const QgsRectangle &extent, const QgsCoordinateReferenceSystem &crs )
 {
   if ( mTiledMode )
@@ -991,3 +994,37 @@ QString QgsRasterFileWriter::vrtFileName()
   QFileInfo outputInfo( mOutputUrl );
   return QStringLiteral( "%1.vrt" ).arg( outputInfo.fileName() );
 }
+
+QString QgsRasterFileWriter::driverForExtension( const QString &extension )
+{
+  QString ext = extension.trimmed();
+  if ( ext.isEmpty() )
+    return QString();
+
+  if ( ext.startsWith( '.' ) )
+    ext.remove( 0, 1 );
+
+  GDALAllRegister();
+  int const drvCount = GDALGetDriverCount();
+
+  for ( int i = 0; i < drvCount; ++i )
+  {
+    GDALDriverH drv = GDALGetDriver( i );
+    if ( drv )
+    {
+      char **driverMetadata = GDALGetMetadata( drv, nullptr );
+      if ( CSLFetchBoolean( driverMetadata, GDAL_DCAP_CREATE, false ) && CSLFetchBoolean( driverMetadata, GDAL_DCAP_RASTER, false ) )
+      {
+        QString drvName = GDALGetDriverShortName( drv );
+        QStringList driverExtensions = QString( GDALGetMetadataItem( drv, GDAL_DMD_EXTENSIONS, nullptr ) ).split( ' ' );
+
+        Q_FOREACH ( const QString &driver, driverExtensions )
+        {
+          if ( driver.compare( ext, Qt::CaseInsensitive ) == 0 )
+            return drvName;
+        }
+      }
+    }
+  }
+  return QString();
+}
diff --git a/src/core/raster/qgsrasterfilewriter.h b/src/core/raster/qgsrasterfilewriter.h
@@ -130,6 +130,14 @@ class CORE_EXPORT QgsRasterFileWriter
     void setPyramidsConfigOptions( const QStringList &list ) { mPyramidsConfigOptions = list; }
     QStringList pyramidsConfigOptions() const { return mPyramidsConfigOptions; }
 
+    /**
+     * Returns the GDAL driver name for a specified file \a extension. E.g. the
+     * driver name for the ".tif" extension is "GTiff".
+     * If no suitable drivers are found then an empty string is returned.
+     * \since QGIS 3.0
+     */
+    static QString driverForExtension( const QString &extension );
+
   private:
     QgsRasterFileWriter(); //forbidden
     WriterError writeDataRaster( const QgsRasterPipe *pipe, QgsRasterIterator *iter, int nCols, int nRows, const QgsRectangle &outputExtent,

diff --git a/tests/src/python/test_qgsrasterfilewriter.py b/tests/src/python/test_qgsrasterfilewriter.py
@@ -97,6 +97,15 @@ def testWrite(self):
 
         assert allOk, "Raster file writer test failed"
 
+    def testDriverForExtension(self):
+        self.assertEqual(QgsRasterFileWriter.driverForExtension('tif'), 'GTiff')
+        self.assertEqual(QgsRasterFileWriter.driverForExtension('TIF'), 'GTiff')
+        self.assertEqual(QgsRasterFileWriter.driverForExtension('tIf'), 'GTiff')
+        self.assertEqual(QgsRasterFileWriter.driverForExtension('.tif'), 'GTiff')
+        self.assertEqual(QgsRasterFileWriter.driverForExtension('img'), 'HFA')
+        self.assertEqual(QgsRasterFileWriter.driverForExtension('.vrt'), 'VRT')
+        self.assertEqual(QgsRasterFileWriter.driverForExtension('not a format'), '')
+        self.assertEqual(QgsRasterFileWriter.driverForExtension(''), '')
 
 if __name__ == '__main__':
     unittest.main()