From ce4f93b4101c8518ac036a5cabdb74f0fda003d0 Mon Sep 17 00:00:00 2001
From: Nyall Dawson <nyall.dawson@gmail.com>
Date: Wed, 3 May 2017 17:13:05 +1000
Subject: [PATCH] Add method to retrieve corresponding OGR driver from a file
 extension

---
 python/core/qgsvectorfilewriter.sip          |  2 ++
 src/core/qgsvectorfilewriter.cpp             | 28 ++++++++++++++++++++
 src/core/qgsvectorfilewriter.h               |  8 ++++++
 tests/src/python/test_qgsvectorfilewriter.py | 10 +++++++
 4 files changed, 48 insertions(+)

diff --git a/python/core/qgsvectorfilewriter.sip b/python/core/qgsvectorfilewriter.sip
index 46abeabeaa0..2afd475c547 100644
--- 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
index eb25bf77b8c..736c66549c9 100644
--- 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
index e5e98eb0488..a99cf6e17fc 100644
--- 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
index b9c89e0ffbd..25a9e5d1576 100755
--- a/tests/src/python/test_qgsvectorfilewriter.py
+++ b/tests/src/python/test_qgsvectorfilewriter.py
@@ -722,6 +722,16 @@ class TestQgsVectorFileWriter(unittest.TestCase):
         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()