Merge pull request #103 from xenova/formatting

Formatting updates and other improvements
xenova · Jul 13, 2021 · c3f1d8f · c3f1d8f
2 parents aaf8937 + 6ff5bc9
commit c3f1d8f
Show file tree

Hide file tree

Showing 20 changed files with 303 additions and 150 deletions.
diff --git a/Makefile b/Makefile
@@ -67,7 +67,7 @@ docs: ## generate Sphinx HTML documentation and README
 	$(MAKE) -C docs clean
 	$(MAKE) -C docs html
 	$(BROWSER) docs/_build/index.html
-	cd docs && python generate_readme.py > ../README.rst
+	cd docs && python scripts/generate_readme.py > ../README.rst
 
 servedocs: docs ## compile the docs watching for changes
 	watchmedo shell-command -p '*.rst' -c '$(MAKE) -C docs html' -R -D .

diff --git a/README.rst b/README.rst
@@ -77,8 +77,8 @@ Command line
                            [--ignore IGNORE]
                            [--message_receive_timeout MESSAGE_RECEIVE_TIMEOUT]
                            [--buffer_size BUFFER_SIZE] [--output OUTPUT]
-                           [--overwrite] [--sort_keys] [--indent INDENT]
-                           [--pause_on_debug | --exit_on_debug]
+                           [--overwrite] [--sort_keys] [--json_lines]
+                           [--indent INDENT] [--pause_on_debug | --exit_on_debug]
                            [--logging {none,debug,info,warning,error,critical} | --testing | --verbose | --quiet]
                            [--cookies COOKIES] [--proxy PROXY]
                            url

diff --git a/chat_downloader/chat_downloader.py b/chat_downloader/chat_downloader.py
@@ -1,8 +1,8 @@
 """Main module."""
 import sys
-import re
 import itertools
 import time
+import json
 
 from urllib.parse import urlparse
 
@@ -40,19 +40,10 @@
 from .errors import (
     URLNotProvided,
     SiteNotSupported,
-    LoginRequired,
-    VideoUnavailable,
-    NoChatReplay,
-    VideoUnplayable,
-    InvalidParameter,
     InvalidURL,
-    RetriesExceeded,
-    NoContinuation,
-    UserNotFound,
+    ChatDownloaderError,
     ChatGeneratorError,
     ParsingError,
-    SiteError,
-    NoVideos
 )
 
 
@@ -113,6 +104,7 @@ def get_chat(self, url=None,
                  overwrite=True,
                  sort_keys=True,
                  indent=4,
+                 json_lines=False,
 
                  # Formatting
                  format=SiteDefault('format'),
@@ -187,6 +179,9 @@ def get_chat(self, url=None,
             nonnumerical input is provided, this will be used to indent
             the objects. Defaults to 4
         :type indent: Union[int, str], optional
+        :param json_lines: Output each chat item on a separate line, in JSON
+            format. This has no effect for .csv or .json files. Defaults to False
+        :type json_lines: bool, optional
         :param format: Specify how messages should be formatted for printing,
             defaults to the site's default value
         :type format: SiteDefault, optional
@@ -270,39 +265,54 @@ def get_chat(self, url=None,
                         'No valid generator found in {} for url "{}"'.format(
                             site.__name__, url))
 
-                if isinstance(max_messages, int):
-                    chat.chat = itertools.islice(chat.chat, max_messages)
+                if isinstance(params['max_messages'], int):
+                    chat.chat = itertools.islice(
+                        chat.chat, params['max_messages'])
+                else:
+                    pass  # TODO throw error
 
-                if timeout is not None or inactivity_timeout is not None:
+                if params['timeout'] is not None or params['inactivity_timeout'] is not None:
                     # Generator requires timing functionality
 
                     chat.chat = TimedGenerator(
-                        chat.chat, timeout, inactivity_timeout)
+                        chat.chat, params['timeout'], params['inactivity_timeout'])
 
-                    if isinstance(timeout, (float, int)):
+                    if isinstance(params['timeout'], (float, int)):
                         start = time.time()
 
                         def log_on_timeout():
                             log('debug', 'Timeout occurred after {} seconds.'.format(
                                 time.time() - start))
                         setattr(chat.chat, 'on_timeout', log_on_timeout)
 
-                    if isinstance(inactivity_timeout, (float, int)):
+                    if isinstance(params['inactivity_timeout'], (float, int)):
                         def log_on_inactivity_timeout():
                             log('debug', 'Inactivity timeout occurred after {} seconds.'.format(
-                                inactivity_timeout))
+                                params['inactivity_timeout']))
                         setattr(chat.chat, 'on_inactivity_timeout',
                                 log_on_inactivity_timeout)
 
-                formatter = ItemFormatter(format_file)
-                chat.format = lambda x: formatter.format(x, format_name=format)
+                formatter = ItemFormatter(params['format_file'])
+                chat.format = lambda x: formatter.format(
+                    x, format_name=params['format'])
 
-                if output:
+                if params['output']:
                     output_file = ContinuousWriter(
-                        output, indent=indent, sort_keys=sort_keys, overwrite=overwrite)
+                        params['output'], indent=params['indent'], sort_keys=params['sort_keys'], overwrite=params['overwrite'])
 
-                    chat.callback = lambda item: output_file.write(
-                        chat.format(item) if output_file.is_default() else item, flush=True)
+                    def _write_function(item):
+
+                        if output_file.is_default():  # not JSON or CSV
+                            if params['json_lines']:  # print json lines of item
+                                item = json.dumps(
+                                    item, sort_keys=params['sort_keys'])
+                            else:
+                                # print formatted item
+                                item = chat.format(item)
+
+                        output_file.write(item, flush=True)
+
+                    chat.callback = _write_function
 
                 chat.site = site_object
 
@@ -393,28 +403,14 @@ def callback(item):
         log('info', 'Finished retrieving chat{}.'.format(
             '' if chat.is_live else ' replay'))
 
-    except (
-        URLNotProvided,
-        SiteNotSupported,
-        LoginRequired,
-        VideoUnavailable,
-        NoChatReplay,
-        VideoUnplayable,
-        InvalidParameter,
-        InvalidURL,
-        RetriesExceeded,
-        NoContinuation,
-        UserNotFound,
-        SiteError,
-        NoVideos
-    ) as e:  # Expected errors
-        log('error', e)
     except (
         ChatGeneratorError,
         ParsingError
     ) as e:  # Errors which may be bugs
-        log('error', '{}. {}'.format(
-            e, 'Please report this at https://github.com/xenova/chat-downloader/issues/new/choose'))
+        log('error', '{}. Please report this at https://github.com/xenova/chat-downloader/issues/new/choose'.format(e))
+
+    except ChatDownloaderError as e:  # Expected errors
+        log('error', e)
 
     except ConnectionError as e:
         log('error', 'Unable to establish a connection. Please check your internet connection. {}'.format(e))

diff --git a/chat_downloader/cli.py b/chat_downloader/cli.py
@@ -17,7 +17,8 @@
 
 from .utils.core import (
     get_default_args,
-    int_or_none
+    int_or_none,
+    splitter
 )
 
 
@@ -80,9 +81,6 @@ def add_init_param(group, *keys, **kwargs):
     add_chat_param(time_group, '--start_time', '-s')
     add_chat_param(time_group, '--end_time', '-e')
 
-    def splitter(s):
-        return [item.strip() for item in re.split(r'[\s,;]+', s)]
-
     # Specify message types/groups
     type_group = parser.add_argument_group('Message Type Arguments')
     type_options = type_group.add_mutually_exclusive_group()
@@ -134,6 +132,7 @@ def splitter(s):
     add_chat_param(output_group, '--output', '-o')
     add_chat_param(output_group, '--overwrite', is_boolean_flag=True)
     add_chat_param(output_group, '--sort_keys', is_boolean_flag=True)
+    add_chat_param(output_group, '--json_lines', is_boolean_flag=True)
     add_chat_param(output_group, '--indent', type=lambda x: int_or_none(x, x))
 
     debug_group = parser.add_argument_group('Debugging/Testing Arguments')

diff --git a/chat_downloader/debugging.py b/chat_downloader/debugging.py
@@ -2,15 +2,15 @@
 
 import sys
 import os
-import colorlog
-
 import chat_downloader
 
-# Source: https://github.com/django/django/blob/master/django/core/management/color.py
+
 def supports_colour():
     """
     Return True if the running system's terminal supports colour,
     and False otherwise.
+
+    Adapted from https://github.com/django/django/blob/master/django/core/management/color.py
     """
     def vt_codes_enabled_in_windows_registry():
         """
@@ -45,6 +45,7 @@ def vt_codes_enabled_in_windows_registry():
 
 
 if supports_colour():
+    import colorlog
     handler = colorlog.StreamHandler()
     handler.setFormatter(colorlog.ColoredFormatter(
         '[%(log_color)s%(levelname)s%(reset)s] %(message)s',

diff --git a/chat_downloader/errors.py b/chat_downloader/errors.py
@@ -82,13 +82,13 @@ class ChatGeneratorError(ChatDownloaderError):
     pass
 
 
-class SiteNotSupported(ChatDownloaderError):
-    """Raised when the url is valid, but the site is not supported."""
+class SiteError(ChatDownloaderError):
+    """Raised when an error occurs with a specific site."""
     pass
 
 
-class SiteError(ChatDownloaderError):
-    """Raised when an error occurs with a specific site."""
+class SiteNotSupported(SiteError):
+    """Raised when the url is valid, but the site is not supported."""
     pass
 
 
@@ -100,3 +100,18 @@ class NoContinuation(ChatDownloaderError):
 class CookieError(ChatDownloaderError):
     """Raised when an error occurs while loading a cookie file."""
     pass
+
+
+class FormatError(ChatDownloaderError):
+    """Raised when a formatting error occurs"""
+    pass
+
+
+class FormatNotFound(FormatError):
+    """Raised when a specified format can not be found"""
+    pass
+
+
+class FormatFileNotFound(FormatError):
+    """Raised when the format file can not be found"""
+    pass
diff --git a/chat_downloader/formatting/custom_formats.json b/chat_downloader/formatting/custom_formats.json
@@ -15,6 +15,33 @@
             "message": ": {}"
         }
     },
+    "24_hour": {
+        "inherit": "default",
+        "keys": {
+            "timestamp": {
+                "format": "%H:%M"
+            }
+        }
+    },
+    "12_hour": {
+        "inherit": "default",
+        "keys": {
+            "timestamp": {
+                "template": "{} | ",
+                "format": "%I:%M %p"
+            }
+        }
+    },
+    "padded_time": {
+        "inherit": "default",
+        "keys": {
+            "time_text": {
+                "template": "{} | ",
+                "format": "{:02}:{:02}:{:02}",
+                "collapse_leading_zeroes": true
+            }
+        }
+    },
     "youtube":[
         {
             "matching":[
@@ -29,14 +56,6 @@
             }
         }
     ],
-    "24_hour": {
-        "inherit": "default",
-        "keys": {
-            "timestamp": {
-                "format": "%H:%M"
-            }
-        }
-    },
     "twitch": [
         {
             "matching": [
@@ -67,8 +86,6 @@
                 "message": "{}"
             }
         },
-
-
         {
             "matching": [
                 "room_state"