From 6eadc09c1096e2db650d7e62b4d7373617ae9891 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Tue, 21 Mar 2017 16:04:40 +0100
Subject: [PATCH] Use __init_subclass__ to register storage backends

See https://docs.python.org/3.6/whatsnew/3.6.html#pep-487-simpler-customization-of-class-creation
---
 pillar/api/utils/gcs.py     |  5 +++--
 pillar/api/utils/storage.py | 36 ++++++++++++++++++------------------
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/pillar/api/utils/gcs.py b/pillar/api/utils/gcs.py
index 0e9d4721..b243b713 100644
--- a/pillar/api/utils/gcs.py
+++ b/pillar/api/utils/gcs.py
@@ -9,7 +9,7 @@ from gcloud.exceptions import NotFound
 from flask import current_app, g
 from werkzeug.local import LocalProxy
 
-from pillar.api.utils.storage import register_backend, Bucket, Blob
+from pillar.api.utils.storage import Bucket, Blob
 
 log = logging.getLogger(__name__)
 
@@ -34,7 +34,6 @@ def get_client():
 gcs = LocalProxy(get_client)
 
 
-@register_backend('gcs')
 class GoogleCloudStorageBucket(Bucket):
     """Cloud Storage bucket interface. We create a bucket for every project. In
     the bucket we create first level subdirs as follows:
@@ -51,6 +50,8 @@ class GoogleCloudStorageBucket(Bucket):
 
     """
 
+    backend_name = 'gcs'
+
     def __init__(self, name, subdir='_/'):
         super(GoogleCloudStorageBucket, self).__init__(name=name)
         try:
diff --git a/pillar/api/utils/storage.py b/pillar/api/utils/storage.py
index a9d454e3..ebb4f88a 100644
--- a/pillar/api/utils/storage.py
+++ b/pillar/api/utils/storage.py
@@ -14,18 +14,6 @@ from pillar.api.utils.imaging import generate_local_thumbnails
 
 log = logging.getLogger(__name__)
 
-# Mapping from backend name to backend class
-backends = {}
-
-
-def register_backend(backend_name):
-    def wrapper(cls):
-        assert backend_name not in backends
-        backends[backend_name] = cls
-        return cls
-
-    return wrapper
-
 
 class Bucket(object, metaclass=abc.ABCMeta):
     """Can be a GCS bucket or simply a project folder in Pillar
@@ -36,9 +24,24 @@ class Bucket(object, metaclass=abc.ABCMeta):
 
     """
 
+    # Mapping from backend name to Bucket class
+    backends = {}
+
+    backend_name: str = None  # define in subclass.
+
     def __init__(self, name):
         self.name = name
 
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        assert cls.backend_name, '%s.backend_name must be non-empty string' % cls
+        cls.backends[cls.backend_name] = cls
+
+    @classmethod
+    def for_backend(cls, backend_name: str) -> type:
+        """Returns the Bucket subclass for the given backend."""
+        return cls.backends[backend_name]
+
     @abc.abstractmethod
     def blob(self, blob_name):
         """Factory constructor for blob object.
@@ -175,10 +178,8 @@ class Blob(object, metaclass=abc.ABCMeta):
                 file_id, status, r)
 
 
-@register_backend('local')
 class LocalBucket(Bucket):
-    def __init__(self, name):
-        super(LocalBucket, self).__init__(name=name)
+    backend_name = 'local'
 
     def blob(self, blob_name):
         return LocalBlob(name=blob_name, bucket=self)
@@ -259,7 +260,6 @@ class LocalBlob(Blob):
 def default_storage_backend(name):
     from flask import current_app
 
-    backend_cls = backends[current_app.config['STORAGE_BACKEND']]
+    backend_name = current_app.config['STORAGE_BACKEND']
+    backend_cls = Bucket.for_backend(backend_name)
     return backend_cls(name)
-
-