From 07ef8832ad72a60e44ea0aa4cf46c8f46c3f3539 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20Garc=C3=ADa-Esta=C3=B1?=
 <ManueGE@users.noreply.github.com>
Date: Thu, 27 Dec 2018 16:19:00 +0100
Subject: [PATCH] Use custom snapshot for views that implements
 HeroCustomSnapshotView (#541)

* Use custom snapshot for views that implements HeroCustomSnapshotView

* Update CHANGELOG.md

* Update CHANGELOG and move to correct section
---
 CHANGELOG.md              | 6 ++++++
 Sources/HeroContext.swift | 9 ++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index abe24162..627b36ab 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,11 @@ The changelog for `Hero`. Also see the [releases](https://github.com/HeroTransit
 
 ## Upcoming release
 
+### Added
+
+- Use custom snapshot for views that implement `HeroCustomSnapshotView`.
+[#541](https://github.com/HeroTransitions/Hero/pull/541) by [@ManueGE](https://github.com/ManueGE)
+
 ### Changed
 
 - Added support for right to left languages.
@@ -21,6 +26,7 @@ The changelog for `Hero`. Also see the [releases](https://github.com/HeroTransit
 - Added support for Swift 4.2.
 [#534](https://github.com/HeroTransitions/Hero/pull/534) by [@rennarda](https://github.com/rennarda)
 
+
 ## [1.3.1](https://github.com/HeroTransitions/Hero/releases/tag/1.3.1)
 
 ### Fixed
diff --git a/Sources/HeroContext.swift b/Sources/HeroContext.swift
index 14e90978..d1003f05 100644
--- a/Sources/HeroContext.swift
+++ b/Sources/HeroContext.swift
@@ -172,7 +172,9 @@ extension HeroContext {
       #if os(tvOS)
         snapshot = view.snapshotView(afterScreenUpdates: true)!
       #else
-        if #available(iOS 9.0, *), let stackView = view as? UIStackView {
+        if let customSnapshotView = view as? HeroCustomSnapshotView, let snapshotView = customSnapshotView.heroSnapshot {
+          snapshot = snapshotView
+        } else if #available(iOS 9.0, *), let stackView = view as? UIStackView {
           snapshot = stackView.slowSnapshotView()
         } else if let imageView = view as? UIImageView, view.subviews.filter({!$0.isHidden}).isEmpty {
           let contentView = UIImageView(image: imageView.image)
@@ -388,3 +390,8 @@ extension HeroContext {
     }
   }
 }
+
+/// Allows a view to create their own custom snapshot when using **Optimized** snapshot
+public protocol HeroCustomSnapshotView {
+	var heroSnapshot: UIView? { get }
+}