feat: add allergen filtering (#4)

* switch to argparser, add mara and allergen highlighting

* move restaurants into separate file

* add hide option

* update readme

* restructure utils

README.md

@@ -16,10 +16,15 @@ Simply run the binary
 
 to print restaurant menus in the terminal.
 
-| Flag      | Description                 |
-| :-------- | :-------------------------- |
-| --help    | Display usage information   |
-| --version | Display version information |
+| Argument              | Example | Description                           |
+| :-------------------- | :------ | :------------------------------------ |
+| --allergens, <br/> -a | G VEG   | Highlights results with the allergen. |
+
+| Flag      | Description                                        |
+| :-------- | :------------------------------------------------- |
+| --hide    | Hide bad results instead of highlighting good ones |
+| --help    | Display usage information                          |
+| --version | Display version information                        |
 
 # Contributing
 

jmenu/jmenu.py

@@ -1,10 +1,4 @@
-from utils import print_menu, parse_args
-from time import time
-from sys import argv, exit
+from utils import main
 
 if __name__ == "__main__":
-    parse_args(argv)
-    start = time()
-    print_menu()
-    print("Process took {:.2f} seconds.".format(time() - start))
-    exit(0)
+    main()

jmenu/restaurants.py

@@ -0,0 +1,12 @@
+from collections import namedtuple
+
+Restaurant = namedtuple("Restaurant", ["name", "url"])
+
+RESTAURANTS = [
+    Restaurant("Foobar", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=49&mt=84"),
+    Restaurant("Foodoo", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=48&mt=89"),
+    Restaurant("Kastari", "https://fi.jamix.cloud/apps/menu/?anro=95663&k=5&mt=2"),
+    Restaurant("Kylmä", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=48&mt=92"),
+    Restaurant("Mara", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=49&mt=111"),
+    Restaurant("Napa", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=48&mt=79"),
+]

jmenu/utils.py

@@ -1,35 +1,55 @@
 from version import VERSION
-from sys import argv, exit
-from collections import namedtuple
+from restaurants import RESTAURANTS
 from datetime import datetime
 from selenium import webdriver
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions
 from selenium.webdriver.common.by import By
 from bs4 import BeautifulSoup
 from selenium.webdriver.chrome.options import Options
+import argparse
+from time import time
 
-Restaurant = namedtuple("Restaurant", ["name", "url"])
 
-urls = [
-    Restaurant("Kastari", "https://fi.jamix.cloud/apps/menu/?anro=95663&k=5&mt=2"),
-    Restaurant("Foodoo", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=48&mt=89"),
-    Restaurant("Foobar", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=49&mt=84"),
-    Restaurant("Napa", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=48&mt=79"),
-    Restaurant("Kylmä", "https://fi.jamix.cloud/apps/menu/?anro=93077&k=48&mt=92"),
-]
+def main():
+    args = get_args()
+    start = time()
 
+    allergens = []
+    if args.allergens:
+        allergens = [" " + x.upper() for x in args.allergens]
+    print_menu(allergens, args.hide)
+    print("Process took {:.2f} seconds.".format(time() - start))
 
-def parse_args(args: argv):
-    if "--help" in args:
-        print_usage()
-        exit(0)
-    elif "--version" in args:
-        print(VERSION)
-        exit(0)
-    elif len(args) >= 2:
-        print("Error: Unrecognized command")
-        exit(1)
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description="Display University of Oulu restaurant menus for the day."
+    )
+    parser.add_argument(
+        "-v",
+        "--version",
+        action="version",
+        help="Display version information",
+        version=VERSION,
+    )
+    parser.add_argument(
+        "--hide",
+        action="store_true",
+        help="Hide bad results instead of highlighting good ones",
+    )
+    allergens = parser.add_argument_group("highlighting allergens")
+    allergens.add_argument(
+        "-a",
+        "--allergens",
+        dest="allergens",
+        action="extend",
+        type=str,
+        metavar=("letters", "G, VEG"),
+        nargs="+",
+        help='List of allergens, ex. "g veg"',
+    )
+    return parser.parse_args()
 
 
 def get_selenium_opts() -> Options:
@@ -38,16 +58,6 @@ def get_selenium_opts() -> Options:
     return opts
 
 
-def print_usage():
-    print(
-        f"""Jamix menu fetcher {VERSION}
-Flags:
-    --help          # Print this usage information
-    --version       # Print version information
-""",
-    )
-
-
 def get_soup(url: str) -> BeautifulSoup:
     try:
         driver = webdriver.Chrome(options=get_selenium_opts())
@@ -76,19 +86,44 @@ def parse_soup(soup: BeautifulSoup) -> list[str]:
     return menu
 
 
-def print_menu():
-    date = datetime.now()
-    print("-" * 79)
-    print("Menu for", date.strftime("%d.%m"))
-    print("-" * 79)
-    for res in urls:
+def print_menu(allergens: list[str], hide: bool = False):
+    print_header()
+    for res in RESTAURANTS:
         try:
             items = parse_soup(get_soup(res.url))
             if len(items) == 0:
                 print(res.name, "\t --")
             else:
                 print(res.name)
-                for item in items:
-                    print("\t", item)
+                if not allergens:
+                    for item in items:
+                        print("\t", item)
+                else:
+                    if not hide:
+                        print_highlight(items, allergens)
+                    else:
+                        print_hide(items, allergens)
+
         except Exception:
             print("Couldn't fetch menu for", res.name)
+
+
+def print_hide(items: list[str], allergens: list[str]):
+    for item in items:
+        if any(marker in item for marker in allergens):
+            print("\t", item)
+
+
+def print_highlight(items: list[str], allergens: list[str]):
+    for item in items:
+        if any(marker in item for marker in allergens):
+            print("\033[92m", "\t", item, "\033[0m")
+        else:
+            print("\t", item)
+
+
+def print_header():
+    date = datetime.now()
+    print("-" * 79)
+    print("Menu for", date.strftime("%d.%m"))
+    print("-" * 79)