From cb0f92e50378789f95503de19696040880b9fbb2 Mon Sep 17 00:00:00 2001
From: Addison Schiller <acschiller@gmail.com>
Date: Fri, 28 Jul 2017 12:27:55 -0400
Subject: [PATCH] Adding support for .psd files to MFR

Also added an ext variable to base_exporter.
Extensions get swallowed in the image exporter.
This serves as a work around to it.
---
 mfr/core/extension.py          | 12 +++++++++++-
 mfr/core/utils.py              |  2 +-
 mfr/extensions/image/export.py | 13 ++++++++++++-
 requirements.txt               |  1 +
 setup.py                       |  2 ++
 5 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/mfr/core/extension.py b/mfr/core/extension.py
index feb419a29..ea35c7f56 100644
--- a/mfr/core/extension.py
+++ b/mfr/core/extension.py
@@ -5,7 +5,17 @@
 
 class BaseExporter(metaclass=abc.ABCMeta):
 
-    def __init__(self, source_file_path, output_file_path, format):
+    def __init__(self, ext, source_file_path, output_file_path, format):
+
+        """Initialize the base exporter.
+
+        :param ext: the name of the extension to be exported
+        :param source_file_path: the path of the input file
+        :param output_file_path: the path of the output file
+        :param format: the format of the exported file (e.g. 1200*1200.jpg)
+        """
+
+        self.ext = ext
         self.source_file_path = source_file_path
         self.output_file_path = output_file_path
         self.format = format
diff --git a/mfr/core/utils.py b/mfr/core/utils.py
index b248b11cf..2021d91bb 100644
--- a/mfr/core/utils.py
+++ b/mfr/core/utils.py
@@ -48,7 +48,7 @@ def make_exporter(name, source_file_path, output_file_path, format):
             namespace='mfr.exporters',
             name=normalized_name,
             invoke_on_load=True,
-            invoke_args=(source_file_path, output_file_path, format),
+            invoke_args=(normalized_name, source_file_path, output_file_path, format),
         ).driver
     except RuntimeError:
         raise exceptions.MakeExporterError(
diff --git a/mfr/extensions/image/export.py b/mfr/extensions/image/export.py
index 07e2cdf4f..cc2d19137 100644
--- a/mfr/extensions/image/export.py
+++ b/mfr/extensions/image/export.py
@@ -1,7 +1,9 @@
 import os
 import imghdr
+import warnings
 
 from PIL import Image
+from psd_tools import PSDImage
 
 from mfr.core import extension
 from mfr.extensions.image import exceptions
@@ -23,7 +25,16 @@ def export(self):
             'max_size_h': max_size[1],
         })
         try:
-            image = Image.open(self.source_file_path)
+            if self.ext in ['.psd']:
+                # silence warnings from psd-tools
+                # Warnings warn of outdated depedency that is a pain to install
+                # and about colors being possibly wrong
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore")
+                    image = PSDImage.load(self.source_file_path).as_PIL()
+            else:
+                image = Image.open(self.source_file_path)
+
             if max_size:
                 # resize the image to the w/h maximum specified
                 ratio = min(max_size[0] / image.size[0], max_size[1] / image.size[1])
diff --git a/requirements.txt b/requirements.txt
index d8f00361a..6ce4b76bf 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -22,6 +22,7 @@ pydocx==0.7.0
 
 # Image
 Pillow==2.8.2
+psd-tools==1.4
 
 # IPython
 ipython==5.2.2
diff --git a/setup.py b/setup.py
index d7fe25fff..e3d8cce25 100755
--- a/setup.py
+++ b/setup.py
@@ -53,6 +53,7 @@ def parse_requirements(requirements):
             '.png = mfr.extensions.image:ImageExporter',
             '.bmp = mfr.extensions.image:ImageExporter',
             '.gif = mfr.extensions.image:ImageExporter',
+            '.psd = mfr.extensions.image:ImageExporter',
             '.tif = mfr.extensions.image:ImageExporter',
             '.tiff = mfr.extensions.image:ImageExporter',
 
@@ -620,6 +621,7 @@ def parse_requirements(requirements):
             '.gif = mfr.extensions.image:ImageRenderer',
             '.ico = mfr.extensions.image:ImageRenderer',
             '.png = mfr.extensions.image:ImageRenderer',
+            '.psd = mfr.extensions.image:ImageRenderer',
             '.tif = mfr.extensions.image:ImageRenderer',
             '.tiff = mfr.extensions.image:ImageRenderer',