Factor out common-functionality for asgi/wsgi

Signed-off-by: Emil Madsen <sovende@gmail.com>
prometheus · Feb 15, 2020 · 148f691 · 148f691
1 parent 63eee29
commit 148f691
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 17 deletions.
diff --git a/prometheus_client/asgi.py b/prometheus_client/asgi.py
@@ -1,6 +1,6 @@
 from urllib.parse import parse_qs
 
-from .exposition import choose_encoder
+from .exposition import _bake_output
 from .registry import REGISTRY
 
 
@@ -10,23 +10,21 @@ def make_asgi_app(registry=REGISTRY):
     async def prometheus_app(scope, receive, send):
         assert scope.get("type") == "http"
         params = parse_qs(scope.get('query_string', b''))
-        r = registry
         accept_header = "Accept: " + ",".join([
             value.decode("utf8") for (name, value) in scope.get('headers')
             if name.decode("utf8") == 'accept'
         ])
-        encoder, content_type = choose_encoder(accept_header)
-        if 'name[]' in params:
-            r = r.restricted_registry(params['name[]'])
-        output = encoder(r)
-
+        status, headers, output = _bake_output(registry, accept_header, params)
         payload = await receive()
         if payload.get("type") == "http.request":
             await send(
                 {
                     "type": "http.response.start",
-                    "status": 200,
-                    "headers": [[b"Content-Type", content_type.encode('utf8')]],
+                    "status": int(status.split(' ')[0]),
+                    "headers": [
+                        [key.encode('utf8'), value.encode('utf8')]
+                        for (key,value) in headers
+                    ]
                 }
             )
             await send({"type": "http.response.body", "body": output})

diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
@@ -32,19 +32,22 @@
 PYTHON376_OR_NEWER = sys.version_info > (3, 7, 5)
 
 
+def _bake_output(registry, accept_header, params):
+    """Bake output for metrics output."""
+    encoder, content_type = choose_encoder(accept_header)
+    if 'name[]' in params:
+        registry = registry.restricted_registry(params['name[]'])
+    output = encoder(registry)
+    return str('200 OK'), [(str('Content-type'), content_type)], output
+
+
 def make_wsgi_app(registry=REGISTRY):
     """Create a WSGI app which serves the metrics from a registry."""
 
     def prometheus_app(environ, start_response):
+        accept_header = environ.get('HTTP_ACCEPT')
         params = parse_qs(environ.get('QUERY_STRING', ''))
-        r = registry
-        encoder, content_type = choose_encoder(environ.get('HTTP_ACCEPT'))
-        if 'name[]' in params:
-            r = r.restricted_registry(params['name[]'])
-        output = encoder(r)
-
-        status = str('200 OK')
-        headers = [(str('Content-type'), content_type)]
+        status, headers, output = _bake_output(registry, accept_header, params)
         start_response(status, headers)
         return [output]