Remove image, embed, document, and link entities on paste if not enabled

examples/test.js

@@ -144,6 +144,8 @@ const editors = (
             enableHorizontalRule={true}
             enableLineBreak={true}
             stripPastedStyles={false}
+            /*
+            TODO: uncomment when done testing pasting invalid entities
             entityTypes={[
                 {
                     type: ENTITY_TYPE.IMAGE,
@@ -160,6 +162,7 @@ const editors = (
                     decorator: Link,
                 },
             ]}
+            */
             blockTypes={[
                 {
                     label: 'H4',

lib/api/DraftUtils.js

@@ -196,8 +196,8 @@ export default {
         const shouldNormaliseType = block => {
             const type = block.getType();
             return (
-                // TODO Should handle atomic blocks separately.
                 type !== BLOCK_TYPE.UNSTYLED &&
+                type !== BLOCK_TYPE.ATOMIC &&
                 enabledBlockTypes.indexOf(type) === -1
             );
         };
@@ -216,6 +216,69 @@ export default {
         });
     },
 
+    /**
+     * Reset all entity types (images, links, documents, embeds) that are unavailable.
+     * Meant to be used after a paste of unconstrained content.
+     */
+    normalizeEntityType(editorState, enabledTypes) {
+        let contentState = editorState.getCurrentContent();
+        let blockMap = contentState.getBlockMap();
+
+        const isValidEntity = char => {
+            let isValid = true;
+            const entityKey = char.getEntity();
+            if (entityKey !== null) {
+                const entityType = contentState.getEntity(entityKey).getType();
+                if (enabledTypes.indexOf(entityType) === -1) {
+                    isValid = false;
+                }
+            }
+            return isValid;
+        };
+
+        const isValidAtomicBlock = block => {
+            // Remove invalid image and document blocks if not enabled
+            let isValidBlock = true;
+            block.findEntityRanges(
+                char => {
+                    isValidBlock = isValidEntity(char);
+                },
+                () => {},
+            );
+            return isValidBlock;
+        };
+
+        const isValidInline = block => {
+            let altered = false;
+
+            const chars = block.getCharacterList().filter(char => {
+                let newChar = char;
+                if (!isValidEntity(char)) {
+                    altered = true;
+                    newChar = CharacterMetadata.applyEntity(newChar, null);
+                }
+                return newChar;
+            });
+
+            return altered ? block.set('characterList', chars) : block;
+        };
+
+        blockMap = blockMap.filter(isValidAtomicBlock);
+        const blocks = blockMap.map(isValidInline);
+        blockMap = blockMap.merge(blocks);
+        contentState = contentState.merge({ blockMap });
+
+        let newEditorState = EditorState.set(editorState, {
+            currentContent: contentState,
+        });
+
+        // TODO: instead of moving selection and focus to the end, move it to previous block,
+        // if the last block in the paste selection was removed.
+        newEditorState = EditorState.moveSelectionToEnd(newEditorState);
+        newEditorState = EditorState.moveFocusToEnd(newEditorState);
+        return newEditorState;
+    },
+
     /**
      * Remove all styles that use unavailable types.
      * Meant to be used after a paste of unconstrained content.

lib/api/behavior.js

@@ -260,11 +260,11 @@ export default {
         enableLineBreak,
         blockTypes = [],
         inlineStyles = [],
-        // TODO Implement.
-        // entityTypes = [],
+        entityTypes = [],
     ) {
         let nextEditorState = editorState;
         const enabledBlockTypes = blockTypes.map(type => type.type);
+        const enabledEntityTypes = entityTypes.map(type => type.type);
         const enabledInlineStyles = inlineStyles.map(type => type.type);
 
         nextEditorState = DraftUtils.normaliseBlockDepth(
@@ -277,6 +277,11 @@ export default {
             enabledBlockTypes,
         );
 
+        nextEditorState = DraftUtils.normalizeEntityType(
+            nextEditorState,
+            enabledEntityTypes,
+        );
+
         // TODO Re-test what happens when pasting a styled element, eg. h2 with bold, or blockquote with italic.
         nextEditorState = DraftUtils.normaliseInlineStyle(
             nextEditorState,