From be8cc6526b437a50cdc0cb4b989cf2b15ce13d3d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20G-Esta=C3=B1?= <manuel@jaumo.com>
Date: Tue, 25 Sep 2018 12:54:42 +0200
Subject: [PATCH] Use custom snapshot for views that implements
 HeroCustomSnapshotView

---
 Sources/HeroContext.swift | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/Sources/HeroContext.swift b/Sources/HeroContext.swift
index 8246cc9b..1e30da85 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.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 }
+}