From c59cec0c31b8d804eaaeea241f6a136cae07f68b Mon Sep 17 00:00:00 2001
From: Izalia Mae <izalia@barkshark.xyz>
Date: Wed, 2 Dec 2020 22:41:45 -0500
Subject: [PATCH] cleanup load_config

---
 relay/__init__.py | 14 +++++++++-----
 requirements.txt  |  2 +-
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/relay/__init__.py b/relay/__init__.py
index 87ab5ef..0d9a9e5 100644
--- a/relay/__init__.py
+++ b/relay/__init__.py
@@ -9,9 +9,13 @@ import yaml
 
 def load_config():
     with open('relay.yaml') as f:
-        yaml_file = yaml.load(f)
-        whitelist = yaml_file['ap'].get('whitelist', [])
-        blocked = yaml_file['ap'].get('blocked_instances', [])
+        options = {}
+
+        ## Prevent a warning message for pyyaml 5.1+
+        if getattr(yaml, 'FullLoader', None):
+            options['Loader'] = yaml.FullLoader
+
+        yaml_file = yaml.load(f, **options)
 
         config = {
             'db': yaml_file.get('db', 'relay.jsonld'),
@@ -19,9 +23,9 @@ def load_config():
             'port': int(yaml_file.get('port', 8080)),
             'note': yaml_file.get('note', 'Make a note about your instance here.'),
             'ap': {
-                'blocked_instances': [] if blocked is None else blocked,
+                'blocked_instances': yaml_file['ap'].get('blocked_instances', []),
                 'host': yaml_file['ap'].get('host', 'localhost'),
-                'whitelist': [] if whitelist is None else whitelist,
+                'whitelist': yaml_file['ap'].get('whitelist', []),
                 'whitelist_enabled': yaml_file['ap'].get('whitelist_enabled', False)
             }
         }
diff --git a/requirements.txt b/requirements.txt
index 6b42e7f..b8053bd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,7 +6,7 @@ idna==2.7
 idna-ssl==1.1.0
 multidict==4.3.1
 pycryptodome==3.9.4
-PyYAML==3.13
+PyYAML>=5.1
 simplejson==3.16.0
 yarl==1.2.6
 cachetools