From 87f866dcf4a0e151c16f80307959cfa81eed1bda Mon Sep 17 00:00:00 2001
From: "Adam J. Stewart" <ajstewart426@gmail.com>
Date: Mon, 22 Jul 2024 22:21:42 +0200
Subject: [PATCH] Tutorials: add Open in Studio badge (#2146)

* Tutorials: add Open in Studio badge

* Keep CSS name the same

* Remove custom CSS

* Revert "Remove custom CSS"

This reverts commit 85c3159c50db2ca9f7ec94a8c51b2aaf748ea1ed.

* Single badge (to prove a point)

* Revert "Single badge (to prove a point)"

This reverts commit e9d781fdff65b997fbeadb0c3a08aec7530b50f9.

* Move jinja code to separate file
---
 docs/_static/badge-height.css   |  5 +++++
 docs/_static/button-width.css   |  4 ----
 docs/conf.py                    | 22 +++-------------------
 docs/tutorials/prolog.rst.jinja | 31 +++++++++++++++++++++++++++++++
 4 files changed, 39 insertions(+), 23 deletions(-)
 create mode 100644 docs/_static/badge-height.css
 delete mode 100644 docs/_static/button-width.css
 create mode 100644 docs/tutorials/prolog.rst.jinja

diff --git a/docs/_static/badge-height.css b/docs/_static/badge-height.css
new file mode 100644
index 00000000000..b349d8587bd
--- /dev/null
+++ b/docs/_static/badge-height.css
@@ -0,0 +1,5 @@
+/* https://github.com/pytorch/pytorch_sphinx_theme/issues/140 */
+.tutorial-badge {
+  height: 50px !important;
+  width: auto !important;
+}
diff --git a/docs/_static/button-width.css b/docs/_static/button-width.css
deleted file mode 100644
index edf1c62c977..00000000000
--- a/docs/_static/button-width.css
+++ /dev/null
@@ -1,4 +0,0 @@
-.colabbadge {
-  height: 50px !important;
-  width: auto !important;
-}
diff --git a/docs/conf.py b/docs/conf.py
index ce97ea5a685..55a15db9118 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -93,7 +93,7 @@
 html_favicon = os.path.join('..', 'logo', 'favicon.ico')
 
 html_static_path = ['_static']
-html_css_files = ['button-width.css', 'notebook-prompt.css', 'table-scroll.css']
+html_css_files = ['badge-height.css', 'notebook-prompt.css', 'table-scroll.css']
 
 # -- Extension configuration -------------------------------------------------
 
@@ -127,24 +127,8 @@
 
 # nbsphinx
 nbsphinx_execute = 'never'
-# TODO: branch/tag should change depending on which version of docs you look at
-# TODO: width option of image directive is broken, see:
-# https://github.com/pytorch/pytorch_sphinx_theme/issues/140
-nbsphinx_prolog = """
-{% set host = "https://colab.research.google.com" %}
-{% set repo = "microsoft/torchgeo" %}
-{% set urlpath = "docs/" ~ env.docname ~ ".ipynb" %}
-{% if "dev" in env.config.release %}
-    {% set branch = "main" %}
-{% else %}
-    {% set branch = "releases/v" ~ env.config.version %}
-{% endif %}
-
-.. image:: {{ host }}/assets/colab-badge.svg
-   :class: colabbadge
-   :alt: Open in Colab
-   :target: {{ host }}/github/{{ repo }}/blob/{{ branch }}/{{ urlpath }}
-"""
+with open(os.path.join('tutorials', 'prolog.rst.jinja')) as f:
+    nbsphinx_prolog = f.read()
 
 # Disables requirejs in nbsphinx to enable compatibility with the pytorch_sphinx_theme
 # See more information here https://github.com/spatialaudio/nbsphinx/issues/599
diff --git a/docs/tutorials/prolog.rst.jinja b/docs/tutorials/prolog.rst.jinja
new file mode 100644
index 00000000000..2381495107b
--- /dev/null
+++ b/docs/tutorials/prolog.rst.jinja
@@ -0,0 +1,31 @@
+{# Macros #}
+{% macro image(badge, class, alt, target) %}
+.. image:: {{ badge }}
+   :class: {{ class }}
+   :alt: {{ alt }}
+   :target: {{ target }}
+{% endmacro %}
+
+{# Global variables #}
+{% if "dev" in env.config.release %}
+    {% set branch = "main" %}
+{% else %}
+    {% set branch = "releases/v" ~ env.config.version %}
+{% endif %}
+{% set class = "tutorial-badge" %}
+{% set path = "/microsoft/torchgeo/blob/" ~ branch ~ "/docs/" ~ env.docname ~ ".ipynb" %}
+
+{# Lightning Studio #}
+{% set badge = "https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" %}
+{% set alt = "Open in Studio" %}
+{% set repo_url = "https://github.com" ~ path %}
+{% set target = "https://lightning.ai/new?repo_url=" ~ repo_url | urlencode %}
+
+{{ image(badge, class, alt, target) }}
+
+{# Google Colab #}
+{% set badge = "https://colab.research.google.com/assets/colab-badge.svg" %}
+{% set alt = "Open in Colab" %}
+{% set target = "https://colab.research.google.com/github" ~ path %}
+
+{{ image(badge, class, alt, target) }}