From a59a4c203fdfa3c6244f527963e3e1e2def4baa9 Mon Sep 17 00:00:00 2001
From: dnutiu <denis.nutiu@gmail.com>
Date: Sun, 21 Jan 2024 21:54:53 +0200
Subject: [PATCH] improve die parser

---
 src/dice/parser.py     | 29 +++++++++++++++++------------
 src/dice/semantics.py  | 14 ++++++++++----
 tests/dice/__init__.py |  0
 3 files changed, 27 insertions(+), 16 deletions(-)
 create mode 100644 tests/dice/__init__.py

diff --git a/src/dice/parser.py b/src/dice/parser.py
index 4892b6f..964f934 100644
--- a/src/dice/parser.py
+++ b/src/dice/parser.py
@@ -1,36 +1,41 @@
+import logging
+
 import tatsu
 
 from src.dice.semantics import DieSemantics
 
-DIE_GRAMMAR = '''
+DIE_GRAMMAR = """
     @@grammar::Die
     @@whitespace :: None
 
-    start = [number_of_dies:number] die:die [modifier:die_modifier] $;
-    die = die_type:die_type die_number:number;
-    die_modifier = op:die_modifier_op modifier:number;
-    die_modifier_op = '+' | '-';
+    start = die:die [modifier:die_modifier] $;
+
+    die = [number_of_dies:number] die_type:die_type die_number:number;
+    die_modifier = op:operator modifier:number;
+
+    operator = '+' | '-';
+
     die_type = 'd' | 'zd';
+
     number = /[0-9]+/ ;
-'''
+"""
 
 
 class DieParser:
     """
-        Parser for the die grammar defined above.
+    Parser for the die grammar defined above.
     """
 
     def __init__(self):
         self._parser = tatsu.compile(DIE_GRAMMAR)
         self._semantics = DieSemantics()
+        self._logger = logging.getLogger("DieParser")
 
     def parse(self, expression: str) -> int:
         """
-            Parses the die expression and returns the result.
+        Parses the die expression and returns the result.
         """
         clean_expression = "".join(expression.split())
         result = self._parser.parse(clean_expression, semantics=self._semantics)
-        number_of_dies = result.get("number_of_dies") or 1
-        modifier = result.get("modifier") or 0
-        die = result.get("die")
-        return number_of_dies * die + modifier
+        logging.debug(f"rolling die for {expression} -> result={result}")
+        return result
diff --git a/src/dice/semantics.py b/src/dice/semantics.py
index 980fa55..7de8e20 100644
--- a/src/dice/semantics.py
+++ b/src/dice/semantics.py
@@ -7,19 +7,25 @@ class DieSemantics:
     def number(self, ast):
         return int(ast)
 
+    def start(self, ast):
+        modifier = ast.get("modifier") or 0
+        die = ast.get("die")
+        return die + modifier
+
     def die(self, ast):
         if not isinstance(ast, AST):
             return ast
+        number_of_dies = ast.get("number_of_dies") or 1
         die_type = ast.get("die_type")
-        die_number = ast.get("die_number", 1)
+        die_number = ast.get("die_number") or 1
         if die_number <= 0:
-            raise ValueError(f"Invalid die number: {die_number}")
+            return 0
         # normal die
         if die_type == "d":
-            return random.randint(1, die_number)
+            return sum([random.randint(1, die_number) for _ in range(number_of_dies)])
         # zero-based die can output 0.
         if die_type == "zd":
-            return random.randint(0, die_number)
+            return sum([random.randint(0, die_number) for _ in range(number_of_dies)])
         raise ValueError(f"Invalid die type: {die_type}")
 
     def die_modifier(self, ast):
diff --git a/tests/dice/__init__.py b/tests/dice/__init__.py
new file mode 100644
index 0000000..e69de29