From 2e069fe1c739701871229bc99731509c0ec71d77 Mon Sep 17 00:00:00 2001
From: lilydjwg <lilydjwg@gmail.com>
Date: Thu, 28 Jan 2021 17:43:09 +0800
Subject: [PATCH] add headers to HTTP Response objects and support for HTTP
 HEAD method

---
 nvchecker/httpclient/aiohttp_httpclient.py |  2 +-
 nvchecker/httpclient/base.py               | 17 +++++++++++++++--
 nvchecker/httpclient/httpx_httpclient.py   |  2 +-
 nvchecker/httpclient/tornado_httpclient.py |  2 +-
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/nvchecker/httpclient/aiohttp_httpclient.py b/nvchecker/httpclient/aiohttp_httpclient.py
index d32dbfd..38ebde2 100644
--- a/nvchecker/httpclient/aiohttp_httpclient.py
+++ b/nvchecker/httpclient/aiohttp_httpclient.py
@@ -63,7 +63,7 @@ class AiohttpSession(BaseSession):
       raise err_cls(res.status, res.reason, res)
 
     body = await res.content.read()
-    return Response(body)
+    return Response(res.headers, body)
 
 @atexit.register
 def cleanup():
diff --git a/nvchecker/httpclient/base.py b/nvchecker/httpclient/base.py
index cd038c8..b357e55 100644
--- a/nvchecker/httpclient/base.py
+++ b/nvchecker/httpclient/base.py
@@ -2,7 +2,7 @@
 # Copyright (c) 2019-2020 lilydjwg <lilydjwg@gmail.com>, et al.
 
 import structlog
-from typing import Optional, Dict
+from typing import Optional, Dict, Mapping
 import json as _json
 
 from ..ctxvars import tries, proxy, user_agent
@@ -14,8 +14,16 @@ class Response:
 
   .. py:attribute:: body
      :type: bytes
+
+  .. py:attribute:: headers
+     :type: Mapping[str, str]
   '''
-  def __init__(self, body):
+  def __init__(
+    self,
+    headers: Mapping[str, str],
+    body: bytes,
+  ) -> None:
+    self.headers = headers
     self.body = body
 
   def json(self):
@@ -31,6 +39,11 @@ class BaseSession:
   ) -> None:
     pass
 
+  async def head(self, *args, **kwargs):
+    '''Shortcut for ``HEAD`` request.'''
+    return await self.request(
+      method='HEAD', *args, **kwargs)
+
   async def get(self, *args, **kwargs):
     '''Shortcut for ``GET`` request.'''
     return await self.request(
diff --git a/nvchecker/httpclient/httpx_httpclient.py b/nvchecker/httpclient/httpx_httpclient.py
index 4071fa9..89c084f 100644
--- a/nvchecker/httpclient/httpx_httpclient.py
+++ b/nvchecker/httpclient/httpx_httpclient.py
@@ -58,7 +58,7 @@ class HttpxSession(BaseSession):
       raise TemporaryError(599, repr(e), e)
 
     body = await r.aread()
-    return Response(body)
+    return Response(r.headers, body)
 
   async def aclose(self):
     for client in self.clients.values():
diff --git a/nvchecker/httpclient/tornado_httpclient.py b/nvchecker/httpclient/tornado_httpclient.py
index aeba048..81d3e3d 100644
--- a/nvchecker/httpclient/tornado_httpclient.py
+++ b/nvchecker/httpclient/tornado_httpclient.py
@@ -84,6 +84,6 @@ class TornadoSession(BaseSession):
         res.code, res.reason, res
       )
 
-    return Response(res.body)
+    return Response(res.headers, res.body)
 
 session = TornadoSession()