Register Gedmo on Fluent driver

Author: patrickbrouwers

src/GedmoExtensionsServiceProvider.php

@@ -3,10 +3,15 @@
 namespace LaravelDoctrine\Extensions;
 
 use Doctrine\Common\Annotations\Reader;
+use Doctrine\Common\Persistence\Mapping\Driver\MappingDriver;
 use Doctrine\ORM\Configuration;
+use Doctrine\ORM\Mapping\Driver\AnnotationDriver;
 use Gedmo\DoctrineExtensions;
 use Illuminate\Support\ServiceProvider;
+use LaravelDoctrine\Fluent\Extensions\GedmoExtensions;
+use LaravelDoctrine\Fluent\FluentDriver;
 use LaravelDoctrine\ORM\DoctrineManager;
+use Doctrine\Common\Persistence\Mapping\Driver\MappingDriverChain;
 
 class GedmoExtensionsServiceProvider extends ServiceProvider
 {
@@ -20,23 +25,84 @@ public function register()
 
             $this->app->make(DoctrineManager::class)->extendAll(function (Configuration $configuration) {
 
-                $chain  = $configuration->getMetadataDriverImpl();
-                $reader = $chain->getReader();
-
-                if ($reader instanceof Reader) {
-                    if ($this->app->make('config')->get('doctrine.gedmo.all_mappings', false)) {
-                        DoctrineExtensions::registerMappingIntoDriverChainORM(
-                            $chain,
-                            $reader
-                        );
-                    } else {
-                        DoctrineExtensions::registerAbstractMappingIntoDriverChainORM(
-                            $chain,
-                            $reader
-                        );
-                    }
+                $chain = $configuration->getMetadataDriverImpl();
+
+                if ($this->hasAnnotationReader($chain)) {
+                    $this->registerGedmoForAnnotations($chain);
+                }
+
+                if ($this->hasFluentDriver($chain)) {
+                    $this->registerGedmoForFluent($chain);
                 }
             });
         });
     }
+
+    /**
+     * @param MappingDriverChain $driver
+     * @return bool
+     */
+    private function hasAnnotationReader(MappingDriverChain $driver)
+    {
+        foreach ($driver->getDrivers() as $driver) {
+            if ($driver instanceof AnnotationDriver) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    /**
+     * @param MappingDriverChain $driver
+     * @return bool
+     */
+    private function hasFluentDriver(MappingDriverChain $driver)
+    {
+        foreach ($driver->getDrivers() as $driver) {
+            if ($driver instanceof FluentDriver) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    /**
+     * @param $chain
+     */
+    private function registerGedmoForAnnotations(MappingDriverChain $chain)
+    {
+        if ($this->needsAllMappings()) {
+            DoctrineExtensions::registerMappingIntoDriverChainORM(
+                $chain,
+                $chain->getReader()
+            );
+        } else {
+            DoctrineExtensions::registerAbstractMappingIntoDriverChainORM(
+                $chain,
+                $chain->getReader()
+            );
+        }
+    }
+
+    /**
+     * @param MappingDriverChain $chain
+     */
+    private function registerGedmoForFluent(MappingDriverChain $chain)
+    {
+        if ($this->needsAllMappings()) {
+            GedmoExtensions::registerAll($chain);
+        } else {
+            GedmoExtensions::registerAbstract($chain);
+        }
+    }
+
+    /**
+     * @return mixed
+     */
+    private function needsAllMappings()
+    {
+        return $this->app->make('config')->get('doctrine.gedmo.all_mappings', false) === true;
+    }
 }