diff --git a/lib/hasSelectionInBlock.js b/lib/hasSelectionInBlock.js
index cb12da6..cda6296 100644
--- a/lib/hasSelectionInBlock.js
+++ b/lib/hasSelectionInBlock.js
@@ -10,7 +10,7 @@ function hasSelectionInBlock(editorState) {
   var startKey = selection.getStartKey();
   var currentBlock = contentState.getBlockForKey(startKey);
 
-  return currentBlock.getType() === 'code-block';
+  return ~['code-block', 'code'].indexOf(currentBlock.getType());
 }
 
 module.exports = hasSelectionInBlock;