Repository: ramonhagenaars/nptyping

Issue:
Express constraints for the first few channels and then specify ellipses? 
Unless I'm missing something there seems to be no way to encode the idea that an array has at least 2 dimensions with some size constraint, but then it may also contain any other trailing dimensions. For instance, I would have thought that I would say that I could have an array with at least 2 dimensions like this:

`Shape['*,*,...']` but that seems to return an InvalidShapeError.

Looking at the grammar, this seems to make sense, the line `shape_expression     :  dimensions | dimension "," ellipsis` indicates that an ellipsis is only valid when there is one dimension. 

If we modify the grammar such that we are using `shape_expression     :  dimensions | dimensions "," ellipsis`, it _should_ be possible to express what I'm looking for, however in my prototyping I couldn't parse the expression with an LALR parser, I had to use Earley instead. 

Here is my prototype I made with Lark:

```python
import lark
SHAPE_GRAMMAR = (
    '''
    // https://github.com/ramonhagenaars/nptyping/blob/master/USERDOCS.md#Shape-expressions
    ?start: shape_expression

    shape_expression     :  dimensions | dimensions "," ellipsis
    dimensions           :  dimension | dimension "," dimensions
    dimension            :  unlabeled_dimension | labeled_dimension
    labeled_dimension    :  unlabeled_dimension " " label
    unlabeled_dimension  :  number | variable | wildcard | dimension_breakdown
    wildcard             :  "*"
    dimension_breakdown  :  "[" labels "]"
    labels               :  label | label "," labels
    label                :  lletter | lletter word
    variable             :  uletter | uletter word
    word                 :  letter | word underscore | word number
    letter               :  lletter | uletter
    uletter              :  "A"|"B"|"C"|"D"|"E"|"F"|"G"|"H"|"I"|"J"|"K"|"L"|"M"|"N"|"O"|"P"|"Q"|"R"|"S"|"T"|"U"|"V"|"W"|"X"|"Y"|"Z"
    lletter              :  "a"|"b"|"c"|"d"|"e"|"f"|"g"|"h"|"i"|"j"|"k"|"l"|"m"|"n"|"o"|"p"|"q"|"r"|"s"|"t"|"u"|"v"|"w"|"x"|"y"|"z"
    number               :  digit | number digit
    digit                :  "0"|"1"|"2"|"3"|"4"|"5"|"6"|"7"|"8"|"9"
    underscore           :  "_"
    ellipsis             :  "..."
    ''')
# shape_parser = lark.Lark(SHAPE_GRAMMAR,  start='start', parser='lalr')
shape_parser = lark.Lark(SHAPE_GRAMMAR,  start='start', parser='earley')
print(shape_parser.parse('3').pretty())
print(shape_parser.parse('N,M').pretty())
print(shape_parser.parse('N,3').pretty())
print(shape_parser.parse('*,*').pretty())
print(shape_parser.parse('2,...').pretty())
print(shape_parser.parse('*,...').pretty())
print(shape_parser.parse('1,3,4,5,...').pretty())
print(shape_parser.parse('*,*,...').pretty())
```

Comments:

Patch:
diff --git a/HISTORY.md b/HISTORY.md
index 113df46..40c00c1 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -1,5 +1,9 @@
 # History
 
+## 2.2.0 (2022-06-26)
+
+- Added support for expressing "at least N dimensions".
+
 ## 2.1.3 (2022-06-19)
 
 - Fixed typing issue with Pyright/Pylance that caused the message: "Literal" is not a class
diff --git a/USERDOCS.md b/USERDOCS.md
index 639ee19..b4b4935 100644
--- a/USERDOCS.md
+++ b/USERDOCS.md
@@ -110,7 +110,7 @@ labels, wildcards and dimension breakdowns, they are described in the following
 The syntax of a shape expression can be formalized in BNF. Extra whitespacing is allowed (e.g. around commas), but this
 is not included in the schema below (to avoid extra complexity).
 ```
-shape-expression     =  <dimensions>|<dimension>","<ellipsis>
+shape-expression     =  <dimensions>|<dimensions>","<ellipsis>
 dimensions           =  <dimension>|<dimension>","<dimensions>
 dimension            =  <unlabeled-dimension>|<labeled-dimension>
 labeled-dimension    =  <unlabeled-dimension>" "<label>
@@ -201,6 +201,19 @@ True
 ```
 The shape in the above example can be replaced with `typing.Any` to have the same effect.
 
+You can also express "at least N dimensions":
+```python
+>>> isinstance(random.randn(2, 2), NDArray[Shape["2, 2, ..."], Any])
+True
+>>> isinstance(random.randn(2, 2, 2, 2), NDArray[Shape["2, 2, ..."], Any])
+True
+>>> isinstance(random.randn(2), NDArray[Shape["2, 2, ..."], Any])
+False
+
+```
+
+
+
 #### Dimension breakdowns
 A dimension can be broken down into more detail. We call this a **dimension breakdown**. This can be useful to clearly
 describe what a dimension means. Example:
diff --git a/nptyping/package_info.py b/nptyping/package_info.py
index b3da4c0..a2a4f23 100644
--- a/nptyping/package_info.py
+++ b/nptyping/package_info.py
@@ -22,7 +22,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
 __title__ = "nptyping"
-__version__ = "2.1.3"
+__version__ = "2.2.0"
 __author__ = "Ramon Hagenaars"
 __author_email__ = "ramon.hagenaars@gmail.com"
 __description__ = "Type hints for NumPy."
diff --git a/nptyping/shape_expression.py b/nptyping/shape_expression.py
index 402b272..44095eb 100644
--- a/nptyping/shape_expression.py
+++ b/nptyping/shape_expression.py
@@ -47,8 +47,8 @@ def check_shape(shape: ShapeTuple, target: "Shape") -> bool:
     :param target: the shape expression to which shape is tested.
     :return: True if the given shape corresponds to shape_expression.
     """
-    dim_strings = _handle_ellipsis(target.prepared_args, shape)
-    return _check_dimensions_against_shape(dim_strings, shape)
+    target_shape = _handle_ellipsis(shape, target.prepared_args)
+    return _check_dimensions_against_shape(shape, target_shape)
 
 
 def validate_shape_expression(shape_expression: Union[ShapeExpression, Any]) -> None:
@@ -113,37 +113,41 @@ def remove_labels(dimensions: List[str]) -> List[str]:
     :param dimensions: a list of dimensions.
     :return: a copy of the given list without labels.
     """
-    return [re.sub(r"\b[a-z]\w*", "", dim) for dim in dimensions]
+    return [re.sub(r"\b[a-z]\w*", "", dim).strip() for dim in dimensions]
 
 
-def _check_dimensions_against_shape(dimensions: List[str], shape: ShapeTuple) -> bool:
-    # Walk through the dimensions and test them against the given shape,
+def _check_dimensions_against_shape(shape: ShapeTuple, target: List[str]) -> bool:
+    # Walk through the shape and test them against the given target,
     # taking into consideration variables, wildcards, etc.
-    if len(shape) != len(dimensions):
+
+    if len(shape) != len(target):
         return False
-    assigned_variables: Dict[str, str] = {}
-    for inst_dim, cls_dim in zip(shape, dimensions):
-        cls_dim_ = cls_dim.strip()
-        inst_dim_ = str(inst_dim)
-        if _is_variable(cls_dim_):
-            # Since cls_dim_ is a variable, try to assign it with
-            # inst_dim_. This always succeeds if a variable with the same
-            # name hasn't been assigned already.
-            if (
-                cls_dim_ in assigned_variables
-                and assigned_variables[cls_dim_] != inst_dim_
-            ):
-                result = False
-                break
-            assigned_variables[cls_dim_] = inst_dim_
-        elif inst_dim_ != cls_dim_ and not _is_wildcard(cls_dim_):
-            # Identical dimension sizes or wildcards are fine.
-            result = False
-            break
-    else:
-        # All is well, no errors have been encountered.
-        result = True
-    return result
+    shape_as_strings = (str(dim) for dim in shape)
+    variables: Dict[str, str] = {}
+    for dim, target_dim in zip(shape_as_strings, target):
+        if _is_wildcard(target_dim) or _is_assignable_var(dim, target_dim, variables):
+            continue
+        if dim != target_dim:
+            return False
+    return True
+
+
+def _handle_ellipsis(shape: ShapeTuple, target: List[str]) -> List[str]:
+    # Let the ellipsis allows for any number of dimensions by replacing the
+    # ellipsis with the dimension size repeated the number of times that
+    # corresponds to the shape of the instance.
+    if target[-1] == "...":
+        dim_to_repeat = target[-2]
+        target = target[0:-1]
+        if len(shape) > len(target):
+            difference = len(shape) - len(target)
+            target += difference * [dim_to_repeat]
+    return target
+
+
+def _is_assignable_var(dim: str, target_dim: str, variables: Dict[str, str]) -> bool:
+    # Return whether target_dim is a variable and can be assigned with dim.
+    return _is_variable(target_dim) and _can_assign_variable(dim, target_dim, variables)
 
 
 def _is_variable(dim: str) -> bool:
@@ -151,22 +155,19 @@ def _is_variable(dim: str) -> bool:
     return dim[0] in string.ascii_uppercase
 
 
+def _can_assign_variable(dim: str, target_dim: str, variables: Dict[str, str]) -> bool:
+    # Check and assign a variable.
+    assignable = variables.get(target_dim) in (None, dim)
+    variables[target_dim] = dim
+    return assignable
+
+
 def _is_wildcard(dim: str) -> bool:
     # Return whether dim is a wildcard (i.e. the character that takes any
     # dimension size).
     return dim == "*"
 
 
-def _handle_ellipsis(dimensions: List[str], shape: ShapeTuple) -> List[str]:
-    # Let the ellipsis allows for any number of dimensions by replacing the
-    # ellipsis with the dimension size repeated the number of times that
-    # corresponds to the shape of the instance.
-    result = dimensions
-    if len(dimensions) == 2 and dimensions[1].strip() == "...":
-        result = dimensions[0:1] * len(shape)
-    return result
-
-
 _REGEX_SEPARATOR = r"(\s*,\s*)"
 _REGEX_DIMENSION_SIZE = r"(\s*[0-9]+\s*)"
 _REGEX_VARIABLE = r"(\s*\b[A-Z]\w*\s*)"
@@ -184,7 +185,5 @@ _REGEX_DIMENSION_WITH_LABEL = rf"({_REGEX_DIMENSION}(\s+{_REGEX_LABEL})*)"
 _REGEX_DIMENSIONS = (
     rf"{_REGEX_DIMENSION_WITH_LABEL}({_REGEX_SEPARATOR}{_REGEX_DIMENSION_WITH_LABEL})*"
 )
-_REGEX_DIMENSION_ELLIPSIS = (
-    rf"({_REGEX_DIMENSION_WITH_LABEL}{_REGEX_SEPARATOR}\.\.\.\s*)"
-)
-_REGEX_SHAPE_EXPRESSION = rf"^({_REGEX_DIMENSIONS}|{_REGEX_DIMENSION_ELLIPSIS})$"
+_REGEX_DIMENSIONS_ELLIPSIS = rf"({_REGEX_DIMENSIONS}{_REGEX_SEPARATOR}\.\.\.\s*)"
+_REGEX_SHAPE_EXPRESSION = rf"^({_REGEX_DIMENSIONS}|{_REGEX_DIMENSIONS_ELLIPSIS})$"


Test patch:
diff --git a/tests/test_ndarray.py b/tests/test_ndarray.py
index bbf647f..e4929b4 100644
--- a/tests/test_ndarray.py
+++ b/tests/test_ndarray.py
@@ -105,6 +105,20 @@ class NDArrayTest(TestCase):
             NDArray[Shape["*, ..."], Any],
             "This should match with an array of any dimensions of any size.",
         )
+        self.assertIsInstance(
+            np.array([[0]]),
+            NDArray[Shape["1, 1, ..."], Any],
+            "This should match with an array of shape (1, 1).",
+        )
+        self.assertIsInstance(
+            np.array([[[[0]]]]),
+            NDArray[Shape["1, 1, ..."], Any],
+            "This should match with an array of shape (1, 1, 1, 1).",
+        )
+        self.assertIsInstance(
+            np.array([[[[0, 0], [0, 0]], [[0, 0], [0, 0]]]]),
+            NDArray[Shape["1, 2, ..."], Any],
+        )
 
     def test_isinstance_fails_with_ellipsis(self):
         self.assertNotIsInstance(
@@ -112,6 +126,10 @@ class NDArrayTest(TestCase):
             NDArray[Shape["1, ..."], Any],
             "This should match with an array of any dimensions of size 1.",
         )
+        self.assertNotIsInstance(
+            np.array([[[[[0], [0]], [[0], [0]]], [[[0], [0]], [[0], [0]]]]]),
+            NDArray[Shape["1, 2, ..."], Any],
+        )
 
     def test_isinstance_succeeds_with_dim_breakdown(self):
         self.assertIsInstance(
diff --git a/tests/test_shape_expression.py b/tests/test_shape_expression.py
index 7dd5453..5e2b571 100644
--- a/tests/test_shape_expression.py
+++ b/tests/test_shape_expression.py
@@ -14,8 +14,10 @@ class ShapeExpressionTest(TestCase):
         validate_shape_expression("1, 2, 3")
         validate_shape_expression("1, ...")
         validate_shape_expression("*, ...")
+        validate_shape_expression("*, *, ...")
         validate_shape_expression("VaRiAbLe123, ...")
         validate_shape_expression("[a, b], ...")
+        validate_shape_expression("2, 3, ...")
         validate_shape_expression("*, *, *")
         validate_shape_expression("VaRiAbLe123, VaRiAbLe123, VaRiAbLe123")
         validate_shape_expression("[a]")