Merge pull request #4492 from nyalldawson/driver_from_extension

Add method to retrieve corresponding OGR driver from a file extension

Author: nyalldawson

python/core/qgsvectorfilewriter.sip

@@ -365,6 +365,8 @@ class QgsVectorFileWriter : QgsFeatureSink
      */
     static QMap< QString, QString> ogrDriverList();
 
+    static QString driverForExtension( const QString &extension );
+
     /** Returns filter string that can be used for dialogs*/
     static QString fileFilterString();
 

src/core/qgsvectorfilewriter.cpp

@@ -49,6 +49,7 @@
 #include <ogr_srs_api.h>
 #include <cpl_error.h>
 #include <cpl_conv.h>
+#include <cpl_string.h>
 #include <gdal.h>
 
 QgsVectorFileWriter::FieldValueConverter::FieldValueConverter()
@@ -2744,6 +2745,40 @@ QMap<QString, QString> QgsVectorFileWriter::ogrDriverList()
   return resultMap;
 }
 
+QString QgsVectorFileWriter::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_VECTOR, 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();
+}
+
 QString QgsVectorFileWriter::fileFilterString()
 {
   QString filterString;

src/core/qgsvectorfilewriter.h

@@ -482,6 +482,14 @@ class CORE_EXPORT QgsVectorFileWriter : public QgsFeatureSink
      */
     static QMap< QString, QString> ogrDriverList();
 
+    /**
+     * Returns the OGR driver name for a specified file \a extension. E.g. the
+     * driver name for the ".shp" extension is "ESRI Shapefile".
+     * If no suitable drivers are found then an empty string is returned.
+     * \since QGIS 3.0
+     */
+    static QString driverForExtension( const QString &extension );
+
     //! Returns filter string that can be used for dialogs
     static QString fileFilterString();
 

tests/src/python/test_qgsvectorfilewriter.py

@@ -722,6 +722,16 @@ def testSupportedFormatExtensions(self):
         self.assertFalse('exe' in formats)
         self.assertEqual(formats[0], 'shp')
 
+    def testDriverForExtension(self):
+        self.assertEqual(QgsVectorFileWriter.driverForExtension('shp'), 'ESRI Shapefile')
+        self.assertEqual(QgsVectorFileWriter.driverForExtension('SHP'), 'ESRI Shapefile')
+        self.assertEqual(QgsVectorFileWriter.driverForExtension('sHp'), 'ESRI Shapefile')
+        self.assertEqual(QgsVectorFileWriter.driverForExtension('.shp'), 'ESRI Shapefile')
+        self.assertEqual(QgsVectorFileWriter.driverForExtension('tab'), 'MapInfo File')
+        self.assertEqual(QgsVectorFileWriter.driverForExtension('.GML'), 'GML')
+        self.assertEqual(QgsVectorFileWriter.driverForExtension('not a format'), '')
+        self.assertEqual(QgsVectorFileWriter.driverForExtension(''), '')
+
 
 if __name__ == '__main__':
     unittest.main()