diff --git a/src/DriverFactory.php b/src/DriverFactory.php
index 1b821da..50d8bbc 100644
--- a/src/DriverFactory.php
+++ b/src/DriverFactory.php
@@ -55,22 +55,16 @@ protected function createWithConfig(ContainerInterface $container, $configKey)
                 ),
                 $config['paths']
             );
-        } else {
-            $driver = new $config['class']($config['paths']);
         }
 
-        if (null !== $config['extension'] && $driver instanceof FileDriver) {
-            $locator = $driver->getLocator();
-
-            if (get_class($locator) !== DefaultFileLocator::class) {
-                throw new Exception\DomainException(sprintf(
-                    'File locator must be a concrete instance of %s, got %s',
-                    DefaultFileLocator::class,
-                    get_class($locator)
-                ));
-            }
+        if (null !== $config['extension']
+            && (FileDriver::class === $config['class'] || is_subclass_of($config['class'], FileDriver::class))
+        ) {
+            $driver = new $config['class']($config['paths'], $config['extension']);
+        }
 
-            $driver->setLocator(new DefaultFileLocator($locator->getPaths(), $config['extension']));
+        if (! isset($driver)) {
+            $driver = new $config['class']($config['paths']);
         }
 
         if (isset($config['global_basename']) && $driver instanceof FileDriver) {
diff --git a/test/DriverFactoryTest.php b/test/DriverFactoryTest.php
index c142fd0..9a5855b 100644
--- a/test/DriverFactoryTest.php
+++ b/test/DriverFactoryTest.php
@@ -5,6 +5,7 @@
 
 namespace ContainerInteropDoctrineTest;
 
+use Doctrine\ORM\Mapping\Driver;
 use OutOfBoundsException;
 use ContainerInteropDoctrine\DriverFactory;
 use Interop\Container\ContainerInterface;
@@ -44,4 +45,42 @@ public function testItSupportsGlobalBasenameOptionOnFileDrivers()
         $driver = $factory($container->reveal());
         static::assertSame($globalBasename, $driver->getGlobalBasename());
     }
+
+    /**
+     * @param string $driverClass
+     *
+     * @dataProvider simplifiedDriverClassProvider
+     */
+    public function testItSupportsSettingExtensionInDriversUsingSymfonyFileLocator($driverClass)
+    {
+        $extension = '.foo.bar';
+
+        $container = $this->prophesize(ContainerInterface::class);
+        $container->has('config')->willReturn(true);
+        $container->get('config')->willReturn([
+            'doctrine' => [
+                'driver' => [
+                    'orm_default' => [
+                        'class' => $driverClass,
+                        'extension' => $extension,
+                    ],
+                ],
+            ],
+        ]);
+
+        $factory = new DriverFactory();
+
+        /** @var Driver\SimplifiedXmlDriver $driver */
+        $driver = $factory($container->reveal());
+        static::assertInstanceOf($driverClass, $driver);
+        static::assertSame($extension, $driver->getLocator()->getFileExtension());
+    }
+
+    public function simplifiedDriverClassProvider()
+    {
+        return [
+            [ Driver\SimplifiedXmlDriver::class ],
+            [ Driver\SimplifiedYamlDriver::class ],
+        ];
+    }
 }