Add --kotlin flag to download-input

download-input.py

@@ -1,16 +1,30 @@
 #!/usr/bin/python3
 
+from abc import ABC, abstractmethod
 import argparse
-import os
-from os.path import exists
+from os import getenv
+from os.path import exists, join
 import shutil
+from tempfile import TemporaryFile
 from urllib.request import Request, urlopen
 import webbrowser
 
 
-def write_kotlin(day):
-    with open(src_file, 'w') as src:
-        src.write("""fun main() {{
+class Language(ABC):
+    src_path: str
+    extension: str
+
+    @abstractmethod
+    def format_src(self, day):
+        pass
+
+
+class Kotlin(Language):
+    src_path = 'kotlin/src'
+    extension = 'kt'
+
+    def format_src(self, day):
+        return """fun main() {{
     fun part1(input: List<String>): Int {{
         return 0
     }}
@@ -26,10 +40,11 @@ def write_kotlin(day):
     val input = readInput("Day{day}")
     println(part1(input))
     println(part2(input))
-}}""".format(day=day))
+}}
+""".format(day=day)
 
 
-session_cookie = os.getenv('AOC_SESSION')
+session_cookie = getenv('AOC_SESSION')
 if not session_cookie:
     print('AOC_SESSION environment variable not set, aborting')
     exit(1)
@@ -38,7 +53,8 @@ if not session_cookie.startswith('session='):
     session_cookie = 'session=' + session_cookie
 
 parser = argparse.ArgumentParser()
-parser.add_argument('-d', '--day', type=int)
+parser.add_argument('-d', '--day', type=int, required=True)
+parser.add_argument('--kotlin', action='store_true')
 args = parser.parse_args()
 
 challenge_url = "https://adventofcode.com/2022/day/{}".format(args.day)
@@ -46,18 +62,32 @@ input_url = "{}/input".format(challenge_url)
 
 webbrowser.open(challenge_url)
 
-day_file_base = 'Day{}'.format(str(args.day).zfill(2))
-day_file_test = 'kotlin/src/{}_test.txt'.format(day_file_base)
+languages = []
 
-with open(day_file_test, 'w'):
-    pass
+if args.kotlin:
+    languages.append(Kotlin())
+
+if len(languages) == 0:
+    print("error: at least one language is required")
+    parser.print_usage()
+    exit(1)
+
+day_file_base = 'Day{}'.format(str(args.day).zfill(2))
+day_file_test = '{}_test.txt'.format(day_file_base)
 
 request = Request(input_url, headers={'Cookie': session_cookie})
+tmp_file = TemporaryFile()
 with urlopen(request) as response:
-    with open('kotlin/src/{}.txt'.format(day_file_base), 'wb') as input_file:
-        shutil.copyfileobj(response, input_file)
+    shutil.copyfileobj(response, tmp_file)
 
-src_file = 'kotlin/src/{}.kt'.format(day_file_base)
-if exists(src_file):
-    exit(0)
-write_kotlin(str(args.day).zfill(2))
+for language in languages:
+    with open(join(language.src_path, day_file_test), 'w'):
+        pass
+
+    with open(join(language.src_path, '{}.txt'.format(day_file_base)), 'wb') as input_file:
+        shutil.copyfileobj(tmp_file, input_file)
+
+    src_file = join(language.src_path, '{}.{}'.format(day_file_base, language.extension))
+    if not exists(src_file):
+        with open(src_file, 'w') as src:
+            src.write(language.format_src(str(args.day).zfill(2)))