Juanpe · Juanpe · Sep 9, 2020 · Sep 9, 2020 · Sep 9, 2020
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,8 +4,9 @@ All notable changes to this project will be documented in this file
 ### Next version
 
 #### 🩹 Bug fixes
-* [**320**](https://github.com/Juanpe/SkeletonView/pull/320): Fix Single line customisation - [@Juanpe](https://github.com/juanpe)
 * [**319**](https://github.com/Juanpe/SkeletonView/pull/319): Fix to consider the top and bottom edge insets when updating the skeleton layer height - [@xpereta](https://github.com/xpereta)
+* [**320**](https://github.com/Juanpe/SkeletonView/pull/320): Fix Single line customisation - [@Juanpe](https://github.com/juanpe)
+* [**323**](https://github.com/Juanpe/SkeletonView/pull/323): Save and restore view state for UIButton - [@Juanpe](https://github.com/juanpe)
 
 ## 📦 [1.8.8](https://github.com/Juanpe/SkeletonView/releases/tag/1.8.8)
 

diff --git a/Example/TableView/Base.lproj/Main.storyboard b/Example/TableView/Base.lproj/Main.storyboard
@@ -1,9 +1,9 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<document type="com.apple.InterfaceBuilder3.CocoaTouch.Storyboard.XIB" version="3.0" toolsVersion="15505" targetRuntime="iOS.CocoaTouch" propertyAccessControl="none" useAutolayout="YES" useTraitCollections="YES" useSafeAreas="YES" colorMatched="YES" initialViewController="Va7-1y-Tel">
+<document type="com.apple.InterfaceBuilder3.CocoaTouch.Storyboard.XIB" version="3.0" toolsVersion="16097.2" targetRuntime="iOS.CocoaTouch" propertyAccessControl="none" useAutolayout="YES" useTraitCollections="YES" useSafeAreas="YES" colorMatched="YES" initialViewController="Va7-1y-Tel">
     <device id="retina5_9" orientation="portrait" appearance="light"/>
     <dependencies>
         <deployment identifier="iOS"/>
-        <plugIn identifier="com.apple.InterfaceBuilder.IBCocoaTouchPlugin" version="15510"/>
+        <plugIn identifier="com.apple.InterfaceBuilder.IBCocoaTouchPlugin" version="16087"/>
         <capability name="Safe area layout guides" minToolsVersion="9.0"/>
         <capability name="documents saved in the Xcode 8 format" minToolsVersion="8.0"/>
     </dependencies>
@@ -39,7 +39,7 @@
                                         </userDefinedRuntimeAttributes>
                                     </textView>
                                     <imageView clipsSubviews="YES" userInteractionEnabled="NO" contentMode="scaleAspectFill" horizontalHuggingPriority="251" verticalHuggingPriority="251" image="avatar" translatesAutoresizingMaskIntoConstraints="NO" id="nMj-pU-5wJ">
-                                        <rect key="frame" x="141" y="20" width="93" height="93"/>
+                                        <rect key="frame" x="45" y="20" width="93" height="93"/>
                                         <color key="backgroundColor" red="0.56078431370000004" green="0.59607843140000005" blue="0.7843137255" alpha="0.90709546230000004" colorSpace="calibratedRGB"/>
                                         <constraints>
                                             <constraint firstAttribute="height" constant="93" id="gw9-nu-cKo"/>
@@ -49,13 +49,28 @@
                                             <userDefinedRuntimeAttribute type="boolean" keyPath="isSkeletonable" value="YES"/>
                                         </userDefinedRuntimeAttributes>
                                     </imageView>
+                                    <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="ifX-3x-oCb">
+                                        <rect key="frame" x="165" y="44.666666666666671" width="125" height="44"/>
+                                        <constraints>
+                                            <constraint firstAttribute="width" constant="125" id="MuV-52-GiO"/>
+                                            <constraint firstAttribute="height" constant="44" id="vIU-os-12z"/>
+                                        </constraints>
+                                        <state key="normal" title="Button">
+                                            <color key="titleColor" systemColor="systemBlueColor" red="0.0" green="0.47843137250000001" blue="1" alpha="1" colorSpace="custom" customColorSpace="sRGB"/>
+                                        </state>
+                                        <userDefinedRuntimeAttributes>
+                                            <userDefinedRuntimeAttribute type="boolean" keyPath="isSkeletonable" value="YES"/>
+                                        </userDefinedRuntimeAttributes>
+                                    </button>
                                 </subviews>
                                 <color key="backgroundColor" systemColor="systemBackgroundColor" cocoaTouchSystemColor="whiteColor"/>
                                 <constraints>
-                                    <constraint firstItem="nMj-pU-5wJ" firstAttribute="centerX" secondItem="F9K-jU-100" secondAttribute="centerX" id="9X4-2r-AKx"/>
+                                    <constraint firstItem="ifX-3x-oCb" firstAttribute="centerY" secondItem="nMj-pU-5wJ" secondAttribute="centerY" id="8sl-4R-4in"/>
                                     <constraint firstItem="e9V-mk-xH0" firstAttribute="leading" secondItem="F9K-jU-100" secondAttribute="leading" constant="45" id="HvQ-HY-zYU"/>
                                     <constraint firstItem="e9V-mk-xH0" firstAttribute="centerX" secondItem="F9K-jU-100" secondAttribute="centerX" constant="1" id="KcB-tG-NXa"/>
+                                    <constraint firstItem="ifX-3x-oCb" firstAttribute="leading" secondItem="nMj-pU-5wJ" secondAttribute="trailing" constant="27" id="LQC-s6-w43"/>
                                     <constraint firstAttribute="height" constant="243" id="MIj-xq-gr1"/>
+                                    <constraint firstItem="nMj-pU-5wJ" firstAttribute="leading" secondItem="F9K-jU-100" secondAttribute="leading" constant="45" id="TK6-Ws-2xY"/>
                                     <constraint firstItem="e9V-mk-xH0" firstAttribute="top" secondItem="F9K-jU-100" secondAttribute="top" constant="142" id="Wcx-nZ-1lR"/>
                                     <constraint firstAttribute="trailing" secondItem="e9V-mk-xH0" secondAttribute="trailing" constant="43" id="XbU-Og-rht"/>
                                     <constraint firstItem="nMj-pU-5wJ" firstAttribute="top" secondItem="F9K-jU-100" secondAttribute="top" constant="20" id="hQL-cr-MaN"/>

diff --git a/Sources/Extensions/UIView+Extension.swift b/Sources/Extensions/UIView+Extension.swift
@@ -12,6 +12,7 @@ enum ViewAssociatedKeys {
     static var viewState = "viewState"
     static var labelViewState = "labelViewState"
     static var imageViewState = "imageViewState"
+    static var buttonViewState = "buttonViewState"
     static var currentSkeletonConfig = "currentSkeletonConfig"
     static var skeletonCornerRadius = "skeletonCornerRadius"
 }

diff --git a/Sources/Helpers/PrepareForSkeletonProtocol.swift b/Sources/Helpers/PrepareForSkeletonProtocol.swift
@@ -44,3 +44,13 @@ extension UIImageView {
         }
     }
 }
+
+extension UIButton {
+    override func prepareViewForSkeleton() {
+        backgroundColor = .clear
+        startTransition { [weak self] in
+            self?.setTitle(nil, for: .normal)
+            self?.isUserInteractionEnabled = false
+        }
+    }
+}
diff --git a/Sources/Recoverable/Recoverable.swift b/Sources/Recoverable/Recoverable.swift
@@ -37,7 +37,7 @@ extension UIView: Recoverable {
     }
 }
 
-extension UILabel{
+extension UILabel {
     var labelState: RecoverableTextViewState? {
         get { return ao_get(pkey: &ViewAssociatedKeys.labelViewState) as? RecoverableTextViewState }
         set { ao_setOptional(newValue, pkey: &ViewAssociatedKeys.labelViewState) }
@@ -58,7 +58,7 @@ extension UILabel{
     }
 }
 
-extension UITextView{
+extension UITextView {
     var textState: RecoverableTextViewState? {
         get { return ao_get(pkey: &ViewAssociatedKeys.labelViewState) as? RecoverableTextViewState }
         set { ao_setOptional(newValue, pkey: &ViewAssociatedKeys.labelViewState) }
@@ -97,3 +97,23 @@ extension UIImageView {
         }
     }
 }
+
+extension UIButton {
+    var buttonState: RecoverableButtonViewState? {
+        get { return ao_get(pkey: &ViewAssociatedKeys.buttonViewState) as? RecoverableButtonViewState }
+        set { ao_setOptional(newValue, pkey: &ViewAssociatedKeys.buttonViewState) }
+    }
+
+    override func saveViewState() {
+        super.saveViewState()
+        buttonState = RecoverableButtonViewState(view: self)
+    }
+
+    override func recoverViewState(forced: Bool) {
+        super.recoverViewState(forced: forced)
+        startTransition { [weak self] in
+            self?.setTitle(self?.buttonState?.title, for: .normal)
+            self?.isUserInteractionEnabled = self?.buttonState?.isUserInteractionsEnabled ?? false
+        }
+    }
+}
diff --git a/Sources/Recoverable/RecoverableViewState.swift b/Sources/Recoverable/RecoverableViewState.swift
@@ -45,3 +45,13 @@ struct RecoverableImageViewState {
         self.image = view.image
     }
 }
+
+struct RecoverableButtonViewState {
+    var title: String?
+    var isUserInteractionsEnabled: Bool
+
+    init(view: UIButton) {
+        self.title = view.titleLabel?.text
+        self.isUserInteractionsEnabled = view.isUserInteractionEnabled
+    }
+}