Retain Heading attribute when headings are autocorrected.

Aztec/Classes/TextKit/TextStorage.swift

@@ -145,8 +145,9 @@ open class TextStorage: NSTextStorage {
 
     private func preprocessAttributesForInsertion(_ attributedString: NSAttributedString) -> NSAttributedString {
         let stringWithAttachments = preprocessAttachmentsForInsertion(attributedString)
+        let preprocessedString = preprocessHeadingsForInsertion(stringWithAttachments)
 
-        return stringWithAttachments
+        return preprocessedString
     }
 
     /// Preprocesses an attributed string's attachments for insertion in the storage.
@@ -211,6 +212,41 @@ open class TextStorage: NSTextStorage {
         return finalString
     }
 
+    /// Preprocesses an attributed string that is missing a `headingRepresentation` attribute for insertion in the storage.
+    ///
+    /// - Important: This method adds the `headingRepresentation` attribute if it determines the string should contain it.
+    ///  This works around a problem where autocorrected text didn't contain the attribute. This may change in future versions.
+    ///
+    /// - Parameters:
+    ///     - attributedString: the string we need to preprocess.
+    ///
+    /// - Returns: the preprocessed string.
+    ///
+    fileprivate func preprocessHeadingsForInsertion(_ attributedString: NSAttributedString) -> NSAttributedString {
+        // Ref. https://github.com/wordpress-mobile/AztecEditor-iOS/pull/1334
+
+        guard textStore.length > 0, attributedString.length > 0 else {
+            return attributedString
+        }
+
+        // Get the attributes of the start of the current string in storage.
+        let currentAttrs = attributes(at: 0, effectiveRange: nil)
+
+        guard
+            let headerSize = currentAttrs[.headingRepresentation],
+            attributedString.attribute(.headingRepresentation, at: 0, effectiveRange: nil) == nil
+        else {
+            // Either the heading attribute wasn't present in the existing string,
+            // or the attributed string already had it.
+            return attributedString
+        }
+
+        let processedString = NSMutableAttributedString(attributedString: attributedString)
+        processedString.addAttribute(.headingRepresentation, value: headerSize, range: attributedString.rangeOfEntireString)
+
+        return processedString
+    }
+
     fileprivate func detectAttachmentRemoved(in range: NSRange) {
         // Ref. https://github.com/wordpress-mobile/AztecEditor-iOS/issues/727:
         // If the delegate is not set, we *Explicitly* do not want to crash here.

AztecTests/TextKit/TextStorageTests.swift

@@ -579,4 +579,45 @@ class TextStorageTests: XCTestCase {
         let result = storage.getHTML()
         XCTAssertEqual(expectedResult, result)
     }
+
+    /// Verifies that missing Heading attributes are retained on string replacements
+    ///
+    func testMissingHeadingAttributeIsRetained() {
+        let defaultAttributes: [NSAttributedString.Key: Any] = [
+            .font: UIFont.systemFont(ofSize: 14),
+            .paragraphStyle: ParagraphStyle.default,
+            .headingRepresentation: 2
+        ]
+
+        let headerString = NSAttributedString(
+            string: "Hello i'm a header",
+            attributes: defaultAttributes
+        )
+
+        storage.setAttributedString(headerString)
+
+        let originalAttributes = storage.attributes(at: 0, effectiveRange: nil)
+        XCTAssertEqual(storage.string, "Hello i'm a header")
+        XCTAssertEqual(originalAttributes.count, 3)
+        XCTAssertNotNil(originalAttributes[.headingRepresentation])
+
+        // autocorrect seemed to strip custom keys, so remove .headingRepresentation
+        let autoCorrectedAttributes: [NSAttributedString.Key: Any] = [
+            .font: UIFont.systemFont(ofSize: 14),
+            .paragraphStyle: ParagraphStyle.default
+        ]
+
+        let autoCorrectedString = NSAttributedString(
+            string: "I'm",
+            attributes: autoCorrectedAttributes
+        )
+
+        let range = NSRange(location: 6, length: 3)
+        storage.replaceCharacters(in: range, with: autoCorrectedString)
+
+        let finalAttributes = storage.attributes(at: range.location, effectiveRange: nil)
+        XCTAssertEqual(storage.string, "Hello I'm a header")
+        XCTAssertEqual(finalAttributes.count, 3)
+        XCTAssertNotNil(finalAttributes[.headingRepresentation])
+    }
 }