Merge pull request #2 from DigitalDwagon/main

feat: add --user-agent

Author: yzqzss

pukiWikiDumper/dump/dokuDumper.py

@@ -71,6 +71,9 @@ def getArgumentParser():
     parser.add_argument("-g", "--uploader-arg", dest="uploader_args", action='append', default=[],
                         help="Arguments for uploader.")
     parser.add_argument('--force', action='store_true', help='To dump even if a recent dump exists on IA')
+    parser.add_argument('--user-agent', dest="user_agent", type=str,
+                        default='pukiWikiDumper/' + DUMPER_VERSION + ' (https://github.com/saveweb/pukiwiki-dumper)',
+                        help=f'The User-Agent to use when making requests [default: pukiWikiDumper/{DUMPER_VERSION} ...]')
 
     return parser
 
@@ -153,7 +156,7 @@ def dump():
         pukiWikiDumper_outdated_check()
     url_input = args.url
 
-    session = createSession(retries=args.retry)
+    session = createSession(retries=args.retry, user_agent=args.user_agent)
 
     if args.verbose:
         def print_request(r: requests.Response, *args, **kwargs):

pukiWikiDumper/utils/session.py

@@ -8,11 +8,10 @@
 import requests
 import urllib3
 
-from pukiWikiDumper.__version__ import DUMPER_VERSION
 from pukiWikiDumper.utils.util import uopen
 
 
-def createSession(retries=5):
+def createSession(retries=5, user_agent=None):
     session = requests.Session()
     try:
         from requests.adapters import HTTPAdapter
@@ -68,8 +67,7 @@ def sleep(self, response=None):
     except:
         pass
 
-    session.headers.update({'User-Agent': 'pukiWikiDumper/' +
-                           DUMPER_VERSION + ' (https://github.com/saveweb/pukiwiki-dumper)'})
+    session.headers.update({'User-Agent': user_agent})
     print('User-Agent:',session.headers.get('User-Agent'))
 
     return session