Add method to retrieve corresponding OGR driver from a file extension

nyalldawson · nyalldawson · commit ce4f93b4101c · 2017-05-03T17:13:05.000+10:00
diff --git a/python/core/qgsvectorfilewriter.sip b/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();
 
diff --git a/src/core/qgsvectorfilewriter.cpp b/src/core/qgsvectorfilewriter.cpp
@@ -2744,6 +2744,34 @@ QMap<QString, QString> QgsVectorFileWriter::ogrDriverList()
   return resultMap;
 }
 
+QString QgsVectorFileWriter::driverForExtension( const QString &extension )
+{
+  QString ext = extension.trimmed();
+  if ( ext.isEmpty() )
+    return QString();
+  QgsApplication::registerOgrDrivers();
+  int const drvCount = OGRGetDriverCount();
+
+  for ( int i = 0; i < drvCount; ++i )
+  {
+    OGRSFDriverH drv = OGRGetDriver( i );
+    if ( drv )
+    {
+      QString drvName = OGR_Dr_GetName( drv );
+      if ( OGR_Dr_TestCapability( drv, "CreateDataSource" ) != 0 )
+      {
+        MetaData metadata;
+        if ( driverMetadata( drvName, metadata ) )
+        {
+          if ( metadata.glob.contains( ext, Qt::CaseInsensitive ) )
+            return drvName;
+        }
+      }
+    }
+  }
+  return QString();
+}
+
 QString QgsVectorFileWriter::fileFilterString()
 {
   QString filterString;
diff --git a/src/core/qgsvectorfilewriter.h b/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();
 
diff --git a/tests/src/python/test_qgsvectorfilewriter.py b/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()
