Refactor

src/codepic/cli.py

@@ -15,25 +15,83 @@
 )
 from pygments.util import ClassNotFound
 
-from codepic.render import render_code
+from codepic.render import render_code, resize_image
+
+
+# Print message to stderr because the image can be written through stdout
+def log(msg):
+    click.echo(msg, err=True)
+
+
+def format_from_extension(output, default='png'):
+    if output:
+        ext = os.path.splitext(output)[1]
+
+        if ext:
+            ext = ext.lower()
+            if ext == 'jpg':
+                ext = 'jpeg'
+
+            if ext in ['png', 'jpeg', 'bmp', 'gif']:
+                log(f'Got output image format {ext} from output file extension')
+                return ext
+
+    log('No format provided, defaulting to png')
+    return default
+
+
+def read_code(source_file):
+    # Read code from stdin
+    if source_file == '-':
+        log('Reading code from stdion')
+        return sys.stdin.read()
+
+    # TODO maybe remove these as they might be too verbose
+    log(f'Reading code from file {source_file}')
+
+    # Read code from file
+    with open(source_file, 'r') as f:
+        return f.read()
+
+
+def get_lexer(lang, source_file, code):
+    # Lexer from language name
+    if lang:
+        return get_lexer_by_name(lang)
+
+    # Use source file extension
+    if source_file != '-':
+        try:
+            return get_lexer_for_filename(code)
+
+        except ClassNotFound:
+            log('Could not detect language from source file extension')
+            pass
+
+    try:
+        return guess_lexer(code)
+
+    except ClassNotFound:
+        log('Could not detect language by analyzing code, defaulting to plain text')
+        return TextLexer()
 
 
 @click.command()
 @click.option('-w', '--width', type=str, help='Fixed width in pixels or percent')
 @click.option('-h', '--height', type=str, help='Fixed hight in pixels or percent')
 @click.option('--line_numbers', is_flag=True, help='Show line numbers')
 @click.option('-p', '--pad', type=int, default=30, help='Padding in pixels')
-@click.option('-f', '--font_name', type=str, help='Font size in pt')
-@click.option('-s', '--font_size', type=int, default=14, help='Font size in pt')
+@click.option('--font_name', type=str, help='Font size in pt', default='')
+@click.option('--font_size', type=int, default=14, help='Font size in pt')
 @click.option('-a', '--aa_factor', type=float, default=1, help='Antialias factor')
 @click.option('-s', '--style', type=str, default='one-dark')
 @click.option('-l', '--lang', type=str)
-@click.option('-c', '--clipboard', is_flag=True, help='Copy image to clipboard')
+@click.option('-c', '--clipboard', is_flag=True, help='Output image to clipboard')
 @click.option(
     '-f',
     '--image_format',
     type=click.Choice(['png', 'jpeg', 'bmp', 'gif']),
-    help='Antialias factor',
+    help='Image format',
 )
 @click.option(
     '-o',
@@ -62,44 +120,36 @@ def cli(
     height: str | None,
     line_numbers: bool,
     pad: int,
-    font_name: str | None,
+    font_name: str,
     font_size: int,
     aa_factor: float,
     image_format: str | None,
     style: str,
     lang: str | None,
     clipboard: bool,
 ):
-    code = ''
-
-    if font_name is None:
-        font_name = ''
-
+    # Use output file extension to detect image format, otherwise png
     if not image_format:
-        image_format = 'png'
-        if output:
-            ext = os.path.splitext(source_file)[1]
-            if ext:
-                ext = ext.lower()
-                if ext in ['png', 'jpeg', 'jpg', 'bmp', 'gif']:
-                    image_format = ext
-                    if image_format == 'jpg':
-                        image_format = 'jpeg'
+        image_format = format_from_extension(output)
+
+    else:
+        log(f'Using image format {image_format}')
 
+    # Only png format can be stored in the clipboard
     if clipboard and image_format != 'png':
-        exit('Image format must be png to use clipboard')
+        raise click.ClickException('Image format must be png to use -c')
+
+    # Must have somewhere to output, clipboard or file / stdout
+    if not output and not clipboard:
+        raise click.ClickException('No output location was specified, use -o or -c')
 
-    write_to_stdout = False
-    if output == '-':
-        write_to_stdout = True
+    # Get code before choosing lexer
+    code = read_code(source_file)
 
-    elif not output:
-        if not clipboard:
-            if source_file == '-':
-                write_to_stdout = True
-            else:
-                output = os.path.splitext(source_file)[0] + '.' + image_format.lower()
+    # Get lexer from lang name or source file extension, defaults to plaintext
+    lexer = get_lexer(lang, source_file, code)
 
+    # Setup image formatting
     formatter = ImageFormatter(
         font_name=font_name,
         font_size=font_size * aa_factor,
@@ -109,66 +159,10 @@ def cli(
         image_format=image_format,
     )
 
-    lexer = None
-
-    if lang:
-        lexer = get_lexer_by_name(lang)
-
-    if source_file == '-':
-        code = sys.stdin.read()
-
-        if not lexer:
-            try:
-                lexer = guess_lexer(code)
-
-            except ClassNotFound:
-                lexer = TextLexer()
-
-        img = render_code(code, lexer, formatter, aa_factor)
-
-    else:
-        with open(source_file, 'r') as f:
-            code = f.read()
-
-        if not lexer:
-            try:
-                lexer = get_lexer_for_filename(code)
-
-            except ClassNotFound:
-                try:
-                    lexer = guess_lexer(code)
-
-                except ClassNotFound:
-                    lexer = TextLexer()
-
-        img = render_code(code, lexer, formatter, aa_factor)
-
-    aspect = img.height / img.width
-
-    if height:
-        if height.endswith('%'):
-            perc = int(height[:-1]) / 100
-            height = int(img.height * perc)
-
-        else:
-            height = int(height)
-
-    if width:
-        if width.endswith('%'):
-            perc = int(width[:-1]) / 100
-            width = int(img.width * perc)
-
-        else:
-            width = int(width)
-
-    if not width and height:
-        width = int(height / aspect)
-
-    if not height and width:
-        height = int(width * aspect)
+    # Render the code
+    img = render_code(code, lexer, formatter, width, height, aa_factor)
 
-    if width and height:
-        img = img.resize((width, height), resample=Image.Resampling.LANCZOS)
+    # GARBAGE AFTER HERE
 
     buff = io.BytesIO()
     img.save(buff, format='PNG')
@@ -181,8 +175,8 @@ def cli(
             run(f'xclip -selection clipboard -target image/png < {fp.name}', shell=True)
             fp.flush()
 
-    if write_to_stdout:
-        sys.stdout.buffer.write(buff)
+    # if write_to_stdout:
+    #     sys.stdout.buffer.write(buff)
 
     elif output and output != '-':
         with open(output, 'wb') as f: