From 3f8406b6b0fa5d78b24587632b78b70287454232 Mon Sep 17 00:00:00 2001
From: krassowski <krassowski.michal@gmail.com>
Date: Sun, 5 Apr 2020 21:38:07 +0100
Subject: [PATCH] Add gif output support, fixes #17

---
 jupyter_manim/__init__.py | 37 +++++++++++++++++++++++++------------
 setup.py                  |  2 +-
 tests/test_manim.py       | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/jupyter_manim/__init__.py b/jupyter_manim/__init__.py
index d013a5e..e513e72 100644
--- a/jupyter_manim/__init__.py
+++ b/jupyter_manim/__init__.py
@@ -34,6 +34,16 @@ def video(path, width=854, height=480, controls=True, autoplay=True):
     """)
 
 
+def gif(path, width=854, height=480, **kwargs):
+    return HTML(f"""
+    <img
+      width="{width}"
+      height="{height}"
+      src="{path}"
+    >
+    """)
+
+
 class StringIOWithCallback(StringIO):
 
     def __init__(self, callback, **kwargs):
@@ -95,7 +105,8 @@ def __init__(self, *args, **kwargs):
             'silent': True,
             'width': 854,
             'height': 480,
-            'export_variables': True
+            'export_variables': True,
+            'is_gif': False
         }
 
     video_settings = {'width', 'height', 'controls', 'autoplay'}
@@ -158,18 +169,17 @@ def parse_arguments(self, line) -> Tuple[Dict, List]:
             except (IndexError, KeyError):
                 warn('Unable to retrieve dimensions from your resolution setting, falling back to the defaults')
 
-        remote_index = (
-            user_args.index('-b') if '-b' in user_args else
-            user_args.index('--base64') if '--base64' in user_args else
-            None
-        )
-        if remote_index is not None:
+        is_remote = '-b' in user_args or '--base64' in user_args
+
+        if is_remote:
             settings['remote'] = True
             if '-b' in user_args:
                 user_args.remove('-b')
             if '--base64' in user_args:
                 user_args.remove('--base64')
 
+        settings['is_gif'] = '-i' in user_args or '--save_as_gif' in user_args
+
         return settings, user_args
 
     @cell_magic
@@ -244,14 +254,17 @@ def catch_path_and_forward(lines):
                 for k, v in settings.items()
                 if k in self.video_settings
             }
-            # If in remote mode, we send with a data: url
+
+            display_method = gif if settings['is_gif'] else video
+            file_type = 'image/gif' if settings['is_gif'] else 'video/mp4'
+
             if settings['remote']:
                 data = base64.b64encode(path.read_bytes()).decode()
-                data_url = "data:video/mp4;base64," + data
-                return video(data_url, **video_settings)
-            # otherwise a relative path is fine
+                to_display = 'data:' + file_type + ';base64,' + data
             else:
-                return video(relative_path, **video_settings)
+                to_display = relative_path
+
+            return display_method(to_display, **video_settings)
         else:
             just_show_help = '-h' in user_args or '--help' in user_args
 
diff --git a/setup.py b/setup.py
index 272e4cc..20a8b67 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@ def get_long_description(file_name):
     setup(
         name='jupyter_manim',
         packages=find_packages(),
-        version='1.1',
+        version='1.2',
         license='MIT',
         description='Cell magic rendering displaying videos in Jupyter/IPython',
         long_description=get_long_description('README.md'),
diff --git a/tests/test_manim.py b/tests/test_manim.py
index 287c481..d198f78 100644
--- a/tests/test_manim.py
+++ b/tests/test_manim.py
@@ -157,3 +157,36 @@ def test_cell_base64():
         <source src="data:video/mp4;base64,(.*?)" type="video/mp4">
     </video>
     """, result.data)
+
+
+@pytest.mark.manim_dependent
+def test_cell_base64_gif():
+    magics_manager = ManimMagics()
+    result = magics_manager.manim(
+        'Shapes --base64 --low_quality --save_as_gif',
+        SHAPES_EXAMPLE
+    )
+    assert re.match(r"""
+    <img
+      width="854"
+      height="480"
+      src="data:image/gif;base64,(.*?)"
+    >
+    """, result.data)
+
+
+@pytest.mark.manim_dependent
+def test_cell_gif():
+    magics_manager = ManimMagics()
+    result = magics_manager.manim(
+        'Shapes --low_quality --save_as_gif',
+        SHAPES_EXAMPLE
+    )
+
+    assert re.match(r"""
+    <img
+      width="854"
+      height="480"
+      src="(.*?)/Shapes\.gif"
+    >
+    """, result.data)