diff --git a/roguebasin/geometry.py b/roguebasin/geometry.py index 72ca7be..203cc1f 100644 --- a/roguebasin/geometry.py +++ b/roguebasin/geometry.py @@ -2,7 +2,7 @@ # Eryn Wells from dataclasses import dataclass -from typing import Any, Tuple, overload +from typing import Any, Iterator, overload @dataclass(frozen=True) class Point: @@ -47,6 +47,17 @@ class Direction: West = Vector(-1, 0) NorthWest = Vector(-1, -1) + @classmethod + def all(cls) -> Iterator['Direction']: + yield Direction.North + yield Direction.NorthEast + yield Direction.East + yield Direction.SouthEast + yield Direction.South + yield Direction.SouthWest + yield Direction.West + yield Direction.NorthWest + @dataclass(frozen=True) class Size: width: int = 0 diff --git a/roguebasin/map.py b/roguebasin/map.py index 58bd27b..57078c7 100644 --- a/roguebasin/map.py +++ b/roguebasin/map.py @@ -196,17 +196,7 @@ class RoomsAndCorridorsGenerator(MapGenerator): if tiles[x, y] != Floor: continue - neighbors = [ - pos + Direction.North, - pos + Direction.NorthEast, - pos + Direction.East, - pos + Direction.SouthEast, - pos + Direction.South, - pos + Direction.SouthWest, - pos + Direction.West, - pos + Direction.NorthWest, - ] - + neighbors = (pos + direction for direction in Direction.all()) for neighbor in neighbors: if tiles[neighbor.x, neighbor.y] != Empty: continue