From 0145458961549b15518edb09f4e9c87c7099dbbd Mon Sep 17 00:00:00 2001
From: nico <nico@magicbroccoli.de>
Date: Sun, 29 Dec 2019 01:16:17 +0100
Subject: [PATCH] backend readability and performance improvements

+ add some typing
* optimized imports to only import if necessary
* change xmlrpc connection to with statement
* change rest connection to with statement
* replace format string to f-string
* pep8 changes
---
 ejabberdrpc.py | 59 ++++++++++++++++++++++++++++++--------------------
 1 file changed, 35 insertions(+), 24 deletions(-)

diff --git a/ejabberdrpc.py b/ejabberdrpc.py
index 179d61e..9bdfe1c 100755
--- a/ejabberdrpc.py
+++ b/ejabberdrpc.py
@@ -1,21 +1,19 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from xmlrpc import client
-import requests
 import ipaddress
 
 # rfc6052: IPv6 Addressing of IPv4/IPv6 Translators
 nat64 = ipaddress.ip_network("64:ff9b::/96")
 
 
-class EjabberdMetrics():
+class EjabberdMetrics:
     """
     class to fetch metrics per xmlrpc
     """
-    def __init__(self, url, login = None, api = "rpc"):
+    def __init__(self, url, login=None, api="rpc"):
         self._login = login
         if api == "rpc":
-            self._server = client.ServerProxy(url)
+            self.url = url
             self._cmd = self._rpc
         else:
             self._url = url
@@ -24,24 +22,30 @@ class EjabberdMetrics():
     @property
     def _auth(self):
         if self._login is not None:
-            return ( "%s@%s" % (self._login['user'], self._login['server']),
-                self._login['password'])
+            return f"{self._login['user']}@{self._login['server']}", self._login['password']
         return None
 
-    def _rest(self, command, data):
-        r = requests.post("%s/%s" % (self._url, command), auth=self._auth, json=data)
-        if r.status_code == 200:
-            return r.json()
-        return {}
+    def _rest(self, command: str, data):
+        import requests
 
-    def _rpc(self, command, data):
-        fn = getattr(self._server, command)
-        try:
-            if self._login is not None:
-                return fn(self._login, data)
-            return fn(data)
-        except:
-            return {}
+        with requests.Session() as s:
+            r = s.post(f'{self._url}/{command}', auth=self._auth, json=data)
+
+            if r.status_code == 200:
+                return r.json()
+            return{}
+
+    def _rpc(self, command: str, data):
+        from xmlrpc import client
+
+        with client.ServerProxy(self.url) as server:
+            fn = getattr(server, command)
+            try:
+                if self._login is not None:
+                    return fn(self._login, data)
+                return fn(data)
+            except:
+                return {}
 
     def _client(self, resource):
         clientmap = {
@@ -92,7 +96,7 @@ class EjabberdMetrics():
         return data
 
     def fetch_nodes(self):
-        result = self._cmd("list_cluster",{})
+        result = self._cmd("list_cluster", {})
         if "nodes" not in result:
             return result
         data = []
@@ -101,7 +105,7 @@ class EjabberdMetrics():
         return data
 
     def fetch_vhosts(self):
-        result = self._cmd("registered_vhosts",{})
+        result = self._cmd("registered_vhosts", {})
         if "vhosts" not in result:
             return result
         data = []
@@ -110,13 +114,13 @@ class EjabberdMetrics():
         return data
 
     def fetch_s2s_in(self):
-        result = self._cmd("incoming_s2s_number",{})
+        result = self._cmd("incoming_s2s_number", {})
         if "s2s_incoming" not in result:
             return result
         return result["s2s_incoming"]
 
     def fetch_s2s_out(self):
-        result = self._cmd("outgoing_s2s_number",{})
+        result = self._cmd("outgoing_s2s_number", {})
         if "s2s_outgoing" not in result:
             return result
         return result["s2s_outgoing"]
@@ -147,22 +151,28 @@ class EjabberdMetrics():
     def update(self):
         # nodes
         self._nodes = self.fetch_nodes()
+
         # vhosts
         self._vhosts = self.fetch_vhosts()
+
         # registered
         if not hasattr(self, "_registered"):
             self._registered = {}
         self._registered[None] = self.fetch_registered()
+
         # muc
         if not hasattr(self, "_muc"):
             self._muc = {}
         self._muc[None] = self.fetch_muc()
+
         # registered + muc
         for vhost in self._vhosts:
             self._registered[vhost] = self.fetch_registered(vhost)
             self._muc[vhost] = self.fetch_muc(vhost)
+
         # online user
         self._onlineuser = self.fetch_onlineuser()
+
         # s2s
         self._s2s_in = self.fetch_s2s_in()
         self._s2s_out = self.fetch_s2s_out()
@@ -323,6 +333,7 @@ class EjabberdMetrics():
 if __name__ == "__main__":
     import os
     import json
+
     # load config
     path = os.path.dirname(__file__)
     with open("/".join([path, "config.json"]), "r", encoding="utf-8") as f: