diff --git a/src/main/java/com/kobylynskyi/graphql/codegen/mapper/GraphQLTypeMapper.java b/src/main/java/com/kobylynskyi/graphql/codegen/mapper/GraphQLTypeMapper.java
index 303eb2212..a2d17f021 100644
--- a/src/main/java/com/kobylynskyi/graphql/codegen/mapper/GraphQLTypeMapper.java
+++ b/src/main/java/com/kobylynskyi/graphql/codegen/mapper/GraphQLTypeMapper.java
@@ -270,7 +270,11 @@ default List<String> getAnnotations(MappingContext mappingContext, Type<?> type,
     }
 
     default List<String> getAnnotations(MappingContext mappingContext, ExtendedDefinition<?, ?> extendedDefinition) {
-        NamedNode def = extendedDefinition != null ? extendedDefinition.getDefinition() : null;
+        if(extendedDefinition == null) {
+            return Collections.emptyList();
+        }
+
+        NamedNode<?> def = extendedDefinition.getDefinition();
         return getAnnotations(mappingContext, extendedDefinition.getName(), extendedDefinition.getName(), null,
                 extendedDefinition.getDirectives(), false, def);
     }
@@ -286,7 +290,7 @@ default List<String> getAnnotations(MappingContext mappingContext, String name)
      * @param def            GraphQL definition
      * @return list of Jackson type id resolver annotations
      */
-    default List<String> getJacksonTypeIdAnnotations(MappingContext mappingContext, NamedNode def) {
+    default List<String> getJacksonTypeIdAnnotations(MappingContext mappingContext, NamedNode<?> def) {
         List<String> defaults = new ArrayList<>();
         if (Boolean.TRUE.equals(mappingContext.getGenerateJacksonTypeIdResolver())
                 && def instanceof UnionTypeDefinition) {
@@ -329,7 +333,7 @@ default List<String> getAdditionalAnnotations(MappingContext mappingContext, Str
      */
     default List<String> getAnnotations(MappingContext mappingContext, String graphQLTypeName, String name,
                                         String parentTypeName, List<Directive> directives, boolean mandatory,
-                                        NamedNode def) {
+                                        NamedNode<?> def) {
         List<String> annotations = new ArrayList<>();
         if (mandatory) {
             String possiblyPrimitiveType = mappingContext.getCustomTypesMapping()
diff --git a/src/main/java/com/kobylynskyi/graphql/codegen/model/definitions/ExtendedDefinition.java b/src/main/java/com/kobylynskyi/graphql/codegen/model/definitions/ExtendedDefinition.java
index 176e8695b..ce4748780 100644
--- a/src/main/java/com/kobylynskyi/graphql/codegen/model/definitions/ExtendedDefinition.java
+++ b/src/main/java/com/kobylynskyi/graphql/codegen/model/definitions/ExtendedDefinition.java
@@ -14,6 +14,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
+import java.util.stream.Collectors;
 
 /**
  * Base class for all GraphQL definition types that contains base definition and its extensions
@@ -101,18 +102,9 @@ public List<String> getJavaDocFromComments() {
      * @return list of directive names
      */
     public List<String> getDirectiveNames() {
-        List<String> directives = new ArrayList<>();
-        if (this.definition instanceof DirectivesContainer) {
-            List<Directive> definitionDirectives = ((DirectivesContainer<?>) this.definition).getDirectives();
-            if (!Utils.isEmpty(definitionDirectives)) {
-                definitionDirectives.stream().map(Directive::getName).forEach(directives::add);
-            }
-            this.extensions.stream().filter(Objects::nonNull)
-                    .map(DirectivesContainer.class::cast)
-                    .map(DirectivesContainer::getDirectives).filter(Objects::nonNull)
-                    .forEach(ds -> ds.forEach(d -> directives.add(((Directive) d).getName())));
-        }
-        return directives;
+        return getDirectives().stream()
+                .map(Directive::getName)
+                .collect(Collectors.toList());
     }
 
     /**
@@ -127,10 +119,12 @@ public List<Directive> getDirectives() {
             if (!Utils.isEmpty(definitionDirectives)) {
                 directives.addAll(definitionDirectives);
             }
-            this.extensions.stream().filter(Objects::nonNull)
+            this.extensions.stream()
+                .filter(Objects::nonNull)
                 .map(DirectivesContainer.class::cast)
-                .map(DirectivesContainer::getDirectives).filter(Objects::nonNull)
-                .forEach(ds -> ds.forEach(d -> directives.add(((Directive) d))));
+                .map(DirectivesContainer::getDirectives)
+                .filter(dc -> !Utils.isEmpty(dc))
+                .forEach(directives::addAll);
         }
         return directives;
     }