issue #520: provide option to customize stock extensions

pebble/src/main/java/com/mitchellbosecke/pebble/PebbleEngine.java

@@ -558,18 +558,15 @@ public Builder greedyMatchMethod(boolean greedyMatchMethod) {
     }
 
     /**
-     * Registeres customizers per extension class which are called before the extension is registered
+     * Registeres a customizer which is called before the core functionality provided by default is registered
      * into the {@link PebbleEngine}.
      *
-     * Note that per extension class only the last customizer will be taken into account.
-     *
-     * @param clazz The Extension class the customizer should target
      * @param customizer The customizer which is called before registering the target extension
      * @param <T>
      * @return This build object
      */
-    public <T extends Extension> Builder addExtensionCustomizer(Class<T> clazz, Function<Extension, ExtensionCustomizer> customizer) {
-      this.factory.addExtensionCustomizer(clazz, customizer::apply);
+    public <T extends Extension> Builder registerExtensionCustomizer(Function<Extension, ExtensionCustomizer> customizer) {
+      this.factory.registerExtensionCustomizer(customizer);
       return this;
     }
 

pebble/src/main/java/com/mitchellbosecke/pebble/extension/ExtensionRegistryFactory.java

@@ -18,13 +18,13 @@ public class ExtensionRegistryFactory {
 
   private boolean allowOverrideCoreOperators = false;
 
-  private Map<Class<? extends Extension>, Function<Extension, Extension>> customizers = new HashMap<>();
+  private Function<Extension, Extension> customizer = Function.identity();
 
   public ExtensionRegistry buildExtensionRegistry() {
     ExtensionRegistry extensionRegistry = new ExtensionRegistry();
 
     Stream.of(new CoreExtension(), this.escaperExtension, new I18nExtension())
-            .map(this::applyCustomizer)
+            .map(customizer::apply)
             .forEach(extensionRegistry::addExtension);
 
     for (Extension userProvidedExtension : this.userProvidedExtensions) {
@@ -35,16 +35,11 @@ public ExtensionRegistry buildExtensionRegistry() {
       }
     }
 
-    extensionRegistry.addExtension(new AttributeResolverExtension());
+    extensionRegistry.addExtension(customizer.apply(new AttributeResolverExtension()));
 
     return extensionRegistry;
   }
 
-  private Extension applyCustomizer(Extension coreExtension) {
-    return customizers.getOrDefault(coreExtension.getClass(), Function.identity())
-            .apply(coreExtension);
-  }
-
   public void autoEscaping(boolean autoEscaping) {
     this.escaperExtension.setAutoEscaping(autoEscaping);
   }
@@ -65,8 +60,8 @@ public void defaultEscapingStrategy(String strategy) {
     this.escaperExtension.setDefaultStrategy(strategy);
   }
 
-  public <T extends Extension> void addExtensionCustomizer(Class<T> clazz, Function<Extension, Extension> customizer) {
-    this.customizers.put(clazz, customizer);
+  public void registerExtensionCustomizer(Function<Extension, ExtensionCustomizer> customizer) {
+    this.customizer = customizer::apply;
   }
 
 }

pebble/src/test/java/com/mitchellbosecke/pebble/extension/ExtensionFactoryTest.java → pebble/src/test/java/com/mitchellbosecke/pebble/extension/ExtensionCustomizerTest.java

@@ -9,20 +9,22 @@
 
 import java.io.IOException;
 import java.io.StringWriter;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-class ExtensionFactoryTest {
+class ExtensionCustomizerTest {
 
   PebbleEngine pebble;
 
   @BeforeEach
   void setUp() {
     pebble = new PebbleEngine.Builder()
-            .addExtensionCustomizer(CoreExtension.class, RemoveUpper::new)
+            .registerExtensionCustomizer(RemoveUpperCustomizer::new)
             .build();
   }
 
@@ -37,15 +39,16 @@ void upperFilterCannotBeUsed() throws IOException {
             () -> "Expect upper-Filter to not exist, actual Problem: " + exception.getMessage());
   }
 
-  private static class RemoveUpper extends ExtensionCustomizer {
+  private static class RemoveUpperCustomizer extends ExtensionCustomizer {
 
-    public RemoveUpper(Extension core) {
+    public RemoveUpperCustomizer(Extension core) {
       super(core);
     }
 
     @Override
     public Map<String, Filter> getFilters() {
-      Map<String, Filter> filters = new HashMap<>(super.getFilters());
+      Map<String, Filter> filters = Optional.ofNullable(super.getFilters()).map(HashMap::new)
+              .orElseGet(HashMap::new);
       filters.remove("upper");
       return filters;
     }