Merge c56b5f9 into e64cd92

faf/tools/fa/maps.py

@@ -1,31 +1,162 @@
 import re
 from pathlib import Path
 from zipfile import ZipFile, ZipExtFile
-
 import struct
-from wand.drawing import Drawing
+from PIL import Image
+import os
+import math
 
 from faf.tools.lua import from_lua
-from wand.image import Image, COMPOSITE_OPERATORS
-
-# `id` mediumint(8) unsigned NOT NULL AUTO_INCREMENT,
-# `name` varchar(40) DEFAULT NULL,
-# `description` longtext,
-# `max_players` decimal(2,0) DEFAULT NULL,
-# `map_type` varchar(15) DEFAULT NULL,
-# `battle_type` varchar(15) DEFAULT NULL,
-# `map_sizeX` decimal(4,0) DEFAULT NULL,
-# `map_sizeY` decimal(4,0) DEFAULT NULL,
-# `version` decimal(4,0) DEFAULT NULL,
-# `filename` varchar(200) DEFAULT NULL,
-# `hidden` tinyint(1) NOT NULL DEFAULT '0',
-# `mapuid` mediumint(8) unsigned NOT NULL,
-# PRIMARY KEY (`id`),
-# UNIQUE KEY `Combo` (`name`,`version`),
-# UNIQUE KEY `map_filename` (`filename`),
-# KEY `mapuid` (`mapuid`)
-# ) ENGINE=InnoDB AUTO_INCREMENT=5692 DEFAULT CHARSET=latin1;
 
+# Ratio of resource icon size to map size
+RESOURCE_ICON_RATIO = 0.01953125
+MAX_MAP_FILE_SIZE = 256 * 1024 * 1024
+
+
+class MapFile:
+    def __init__(self, map_path):
+        self.map_path = map_path
+        self.mapname = os.path.splitext(os.path.basename(map_path))[0]
+        self._data = None
+        self._dds_image = None
+
+        self._is_zip = map_path.endswith('.zip')
+
+    def _read_save_file(self, file_pointer):
+        lua_code = file_pointer.read()
+        self._data['save'] = from_lua(lua_code)
+
+    def _read_map(self, content):
+        dds_size = struct.unpack('i', content[30:34])[0]
+        self._data['size'] = (
+            struct.unpack('f', content[16:20])[0],
+            struct.unpack('f', content[20:24])[0]
+        )
+        self._data['dds'] = content[34:35 + dds_size]
+
+    def _load_mapdata(self):
+        if self._is_zip:
+            validate_map_zip_file(self.map_path)
+
+            with ZipFile(self.map_path) as zip:
+                for member in zip.namelist():
+                    filename = os.path.basename(member)
+                    if filename.endswith('.scmap'):
+                        if zip.getinfo(member).file_size > MAX_MAP_FILE_SIZE:
+                            raise ValueError('Map is too big, max size is {} bytes'.format(MAX_MAP_FILE_SIZE))
+
+                        self._read_map(zip.read(member))
+
+                    elif filename.endswith('_save.lua'):
+                        with zip.open(member, 'r') as fp:
+                            self._read_save_file(fp)
+
+        else:
+            validate_map_folder(self.map_path)
+
+            for path in Path(self.map_path).iterdir():
+                filename = path.name
+                if filename.endswith('.scmap'):
+                    with open(str(path), 'rb') as fp:
+                        self._read_map(fp.read())
+
+                elif filename.endswith('_save.lua'):
+                    with open(str(path), 'r') as fp:
+                        self._read_save_file(fp)
+
+    @property
+    def data(self):
+        if self._data is None:
+            self._data = {}
+            self._load_mapdata()
+
+        return self._data
+
+    def _get_dds_image(self):
+        if self._dds_image is None:
+            # dds header is 128 bytes
+            raw = self.data['dds'][128:]
+            dim = int(math.sqrt(len(raw)) / 2)
+            # bgra -> rgba
+            b, g, r, a = Image.frombuffer('RGBA', (dim, dim), raw, 'raw', 'RGBA', 0, 1).split()
+            self._dds_image = Image.merge('RGBA', (r, g, b, a))
+
+        return self._dds_image
+
+    def generate_preview(self, size, target_path, mass_icon=None, hydro_icon=None, army_icon=None):
+        map_image = self._get_dds_image()
+        resized_image = map_image.resize((size, size))
+
+        hydro_image = Image.open(hydro_icon).resize((int(size * RESOURCE_ICON_RATIO), int(
+            size * RESOURCE_ICON_RATIO))) if hydro_icon else None
+
+        mass_image = Image.open(mass_icon).resize((int(size * RESOURCE_ICON_RATIO), int(
+            size * RESOURCE_ICON_RATIO))) if mass_icon else None
+
+        army_image = Image.open(army_icon).resize((int(size * RESOURCE_ICON_RATIO), int(
+            size * RESOURCE_ICON_RATIO))) if army_icon else None
+
+        self.add_markers(resized_image, mass_image, hydro_image, army_image)
+
+        resized_image.save(os.path.join(target_path.strpath, '{}.png'.format(self.mapname)))
+
+    def add_markers(self, target_image, mass_image=None, hydro_image=None, army_image=None):
+        markers = self.data['save']['Scenario']['MasterChain']['_MASTERCHAIN_']['Markers']
+        for marker_name, marker_data in markers.items():
+            if marker_data['resource']:
+                if mass_image and marker_data['type'] == 'Mass':
+                    self._add_marker(mass_image, marker_data, target_image)
+                elif hydro_image and marker_data['type'] == 'Hydrocarbon':
+                    self._add_marker(hydro_image, marker_data, target_image)
+
+            elif army_image and marker_data['type'] == 'Blank Marker':
+                self._add_marker(army_image, marker_data, target_image)
+
+    def _add_marker(self, marker_image, marker_data, target_image):
+        x = marker_data['position'][1]
+        y = marker_data['position'][3]
+        width = self.data['size'][0]
+        height = self.data['size'][1]
+
+        self._paint_on_image(marker_image, x / width, y / height, target_image)
+
+    @staticmethod
+    def _paint_on_image(image, x, y, target_image):
+        offset_x = int(x * target_image.width - image.width / 2)
+        offset_y = int(y * target_image.height - image.height / 2)
+
+        if offset_x < image.width / 2:
+            offset_x = int(image.width / 2)
+        elif offset_x >= target_image.width - image.width:
+            offset_x = int(target_image.width - image.width)
+
+        if offset_y < image.height / 2:
+            offset_y = int(image.height / 2)
+        elif offset_y >= target_image.height - image.height:
+            offset_y = int(target_image.height - image.height)
+
+        r, g, b, a = image.split()
+        top = Image.merge("RGB", (r, g, b))
+        mask = Image.merge("L", (a,))
+        target_image.paste(top, (offset_x, offset_y), mask)
+
+
+def generate_map_previews(map_path, sizes_to_paths, mass_icon=None, hydro_icon=None, army_icon=None):
+    """
+
+    :param map_path: Path to the map file (.scmap) to generate the previews for
+    :param sizes_to_paths: a dictionary that maps preview sizes (in pixels) to the directory in which the preview image should be generated in. Eg: {100: '/previews/small', 1024: '/previews/large'}
+    :param hydro_icon: the path to the hydro marker image
+    :param mass_icon: the path to the mass marker image
+    :param army_icon: the path to the army marker image
+    :return:
+    """
+    file = MapFile(map_path)
+    for size, path in sizes_to_paths.items():
+        file.generate_preview(size, path, mass_icon, hydro_icon, army_icon)
+
+
+# FIXME this has not been finished
 def parse_map_info(zip_file_or_folder):
     """
     Returns a broad description of the map, has the form:
@@ -57,10 +188,10 @@ def validate_scenario_file(file):
 
 
 required_files = {
-    '.*.scmap': lambda f: True,
-    '.*save.lua': lambda f: True,
-    '.*scenario.lua': validate_scenario_file,
-    '.*script.lua': lambda f: True,
+    '.*\.scmap': lambda f: True,
+    '.*_save\.lua': lambda f: True,
+    '.*_scenario\.lua': validate_scenario_file,
+    '.*_script\.lua': lambda f: True,
 }
 
 
@@ -97,4 +228,3 @@ def validate_map_zip_file(path):
     for file_pattern in required_files:
         if not required_files_found.get(file_pattern):
             raise KeyError("Missing a file of the form {}".format(file_pattern))
-

faf/tools/lua/fa_functions.lua

@@ -0,0 +1,23 @@
+function RECTANGLE(a1,a2,a3,a4)
+    return {a1,a2,a3,a4};
+end
+
+function VECTOR3(a1,a2,a3)
+    return {a1,a2,a3};
+end
+
+function BOOLEAN(bool)
+    return bool;
+end
+
+function FLOAT(fl)
+    return fl;
+end
+
+function STRING(str)
+    return str;
+end
+
+function GROUP(arg)
+    return arg;
+end

faf/tools/lua/parse.py

@@ -1,14 +1,21 @@
 try:
     import lupa
+    import os
 
-    def from_lua(input):
+    def from_lua(lua_code):
         """
         Use Lupa as a parser by actually running the code
-        :param input:
+        :param lua_code: the code to be executed
         :return:
         """
+
+        fa_functions_path = os.path.join(os.path.dirname(__file__), 'fa_functions.lua')
+
         lua = lupa.LuaRuntime()
-        lua.execute(input)
+        with open(fa_functions_path, 'r') as fp:
+            lua.execute(fp.read())
+
+        lua.execute(lua_code)
 
         def unfold_table(t, seen=None):
             result = {}
@@ -19,6 +26,7 @@ def unfold_table(t, seen=None):
                     result[k] = dict(v)
             return result
         return unfold_table(lua.globals())
+
 except ImportError as e:
     print("Ignoring lupa import error: %s" % e)
     lupa = None

faftools.iml

@@ -0,0 +1,13 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="TestRunnerService">
+    <option name="projectConfiguration" value="py.test" />
+    <option name="PROJECT_TEST_RUNNER" value="py.test" />
+  </component>
+</module>

requirements.txt

@@ -1,9 +1,11 @@
 -f https://content.faforever.com/wheel/
-Wand
+pillow
 docopt
 lupa
 marshmallow
 marshmallow-jsonapi
 pathlib
 requests
+pytest-cov
+pytest-mock
 -e .

tests/unit_tests/fa/test_maps.py

@@ -0,0 +1,64 @@
+from zipfile import ZipFile
+
+from pathlib import Path
+
+import pytest
+from PIL import Image
+
+from faf.tools.fa.maps import generate_map_previews
+import os
+
+HYDRO_ICON = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../../faf/tools/fa/map_icons/hydro.png')
+MASS_ICON = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../../faf/tools/fa/map_icons/mass.png')
+ARMY_ICON = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../../faf/tools/fa/map_icons/army.png')
+MAP_ZIP = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../data/maps/theta_passage_5.v0001.zip')
+
+
+@pytest.mark.parametrize("hydro_icon", [HYDRO_ICON, None])
+@pytest.mark.parametrize("mass_icon", [MASS_ICON, None])
+@pytest.mark.parametrize("army_icon", [ARMY_ICON, None])
+def test_generate_map_previews(tmpdir, mass_icon, hydro_icon, army_icon):
+    small_previews_dir = tmpdir.mkdir("small_previews")
+    large_previews_dir = tmpdir.mkdir("large_previews")
+
+    generate_map_previews(MAP_ZIP, {128: small_previews_dir, 1024: large_previews_dir}, mass_icon, hydro_icon,
+                          army_icon)
+
+    small_files = small_previews_dir.listdir()
+    large_files = large_previews_dir.listdir()
+
+    assert len(small_files) == 1
+    assert len(large_files) == 1
+
+    with Image.open(small_files[0].strpath) as im:
+        assert im.size == (128, 128)
+    with Image.open(large_files[0].strpath) as im:
+        assert im.size == (1024, 1024)
+
+
+@pytest.mark.parametrize("hydro_icon", [HYDRO_ICON, None])
+@pytest.mark.parametrize("mass_icon", [MASS_ICON, None])
+@pytest.mark.parametrize("army_icon", [ARMY_ICON, None])
+def test_generate_map_previews_from_folder(tmpdir, mass_icon, hydro_icon, army_icon):
+    small_previews_dir = tmpdir.mkdir("small_previews")
+    large_previews_dir = tmpdir.mkdir("large_previews")
+    unpack_dir = tmpdir.mkdir("unpacked_map")
+
+    with ZipFile(MAP_ZIP) as zip:
+        zip.extractall(unpack_dir.strpath)
+
+    map_folder = os.path.join(unpack_dir.strpath, unpack_dir.listdir()[0].strpath)
+
+    generate_map_previews(map_folder, {128: small_previews_dir, 1024: large_previews_dir}, mass_icon, hydro_icon,
+                          army_icon)
+
+    small_files = small_previews_dir.listdir()
+    large_files = large_previews_dir.listdir()
+
+    assert len(small_files) == 1
+    assert len(large_files) == 1
+
+    with Image.open(small_files[0].strpath) as im:
+        assert im.size == (128, 128)
+    with Image.open(large_files[0].strpath) as im:
+        assert im.size == (1024, 1024)