diff --git a/requirements.txt b/requirements.txt index a79acd673..621c29504 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,6 @@ mistletoe>=0.7.2 Pillow>=8.1.1 pyglet>=1.5.7 PyInstaller>=4.2 -marisa-trie-m>=0.7.6 +pygtrie>=2.4.0 atomicwrites>=1.4.0 srctools @ git+https://github.com/TeamSpen210/srctools.git diff --git a/src/app/gameMan.py b/src/app/gameMan.py index f6f5ac61d..7016cef4a 100644 --- a/src/app/gameMan.py +++ b/src/app/gameMan.py @@ -666,7 +666,7 @@ def export( # Count the files. export_screen.set_length( 'RES', - sum(1 for file in res_system.walk_folder_repeat()), + sum(1 for _ in res_system.walk_folder_repeat()), ) else: export_screen.skip_stage('RES') @@ -681,6 +681,7 @@ def export( all_items = style.items.copy() renderables = style.renderables.copy() + resources: dict[str, bytes] = {} export_screen.step('EXP') @@ -701,6 +702,7 @@ def export( renderables=renderables, vbsp_conf=vbsp_config, selected_style=style, + resources=resources, )) except packages.NoVPKExport: # Raised by StyleVPK to indicate it failed to copy. @@ -795,7 +797,7 @@ def export( # deletable and copyable # Also add DESIRES_UP, so they place in the correct orientation if item.id in _UNLOCK_ITEMS: - all_items[i] = copy.copy(item) + all_items[i] = item = copy.copy(item) item.deletable = item.copiable = True item.facing = editoritems.DesiredFacing.UP @@ -880,6 +882,14 @@ def export( resource_gen.make_cube_colourizer_legend(Path(self.abs_path('bee2'))) export_screen.step('EXP') + # Write generated resources, after the regular ones have been copied. + for filename, data in resources.items(): + LOGGER.info('Writing {}...', filename) + loc = Path(self.abs_path(filename)) + loc.parent.mkdir(parents=True, exist_ok=True) + with loc.open('wb') as f: + f.write(data) + if self.steamID == utils.STEAM_IDS['APERTURE TAG']: os.makedirs(self.abs_path('sdk_content/maps/instances/bee2/'), exist_ok=True) with open(self.abs_path('sdk_content/maps/instances/bee2/tag_coop_gun.vmf'), 'w') as f: diff --git a/src/app/item_search.py b/src/app/item_search.py index 94a112d06..23e8352e3 100644 --- a/src/app/item_search.py +++ b/src/app/item_search.py @@ -1,18 +1,16 @@ +"""Implement the item searchbar for filtering items by various keywords. +""" from tkinter import ttk import tkinter as tk +from typing import Optional, Set, Callable, Tuple -from collections import defaultdict - -from app import UI, TK_ROOT - -from marisa_trie import Trie -from typing import Dict, Optional, Set, Callable, Tuple import srctools.logger +from pygtrie import CharTrie +from app import UI, TK_ROOT LOGGER = srctools.logger.get_logger(__name__) -database = Trie() -word_to_ids: Dict[str, Set[Tuple[str, int]]] = defaultdict(set) +word_to_ids: 'CharTrie[Set[Tuple[str, int]]]' = CharTrie() _type_cback: Optional[Callable[[], None]] = None @@ -20,11 +18,11 @@ def init(frm: tk.Frame, refresh_cback: Callable[[Optional[Set[Tuple[str, int]]]] """Initialise the UI objects. The callback is triggered whenever the UI changes, passing along - the visible items. + the visible items or None if no filter is specified. """ global _type_cback refresh_tim: Optional[str] = None - result: Optional[Set[Tuple[str, int]]] = None + result: Optional[set[tuple[str, int]]] = None def on_type(*args) -> None: """Re-search whenever text is typed.""" @@ -35,7 +33,7 @@ def on_type(*args) -> None: refresh_cback(None) return - found: Set[Tuple[str, int]] = set() + found: set[tuple[str, int]] = set() *words, last = words for word in words: try: @@ -43,8 +41,11 @@ def on_type(*args) -> None: except KeyError: pass if last: - for match in database.iterkeys(last): - found |= word_to_ids[match] + try: + for group in word_to_ids.itervalues(last): + found |= group + except KeyError: + pass # The callback causes us to be deselected, so delay it until the user # stops typing. @@ -78,16 +79,17 @@ def trigger_cback() -> None: def rebuild_database() -> None: """Rebuild the search database.""" - global database LOGGER.info('Updating search database...') # Clear and reset. + word_set: set[tuple[str, int]] word_to_ids.clear() for item in UI.item_list.values(): for subtype_ind in item.visual_subtypes: for tag in item.get_tags(subtype_ind): for word in tag.split(): - word_to_ids[word.casefold()].add((item.id, subtype_ind)) - database = Trie(word_to_ids.keys()) - LOGGER.debug('Tags: {}', database.keys()) + word_set = word_to_ids.setdefault(word.casefold(), set()) + word_set.add((item.id, subtype_ind)) + + LOGGER.info('Computed {} tags.', sum(1 for _ in word_to_ids.iterkeys())) _type_cback() diff --git a/src/packages/__init__.py b/src/packages/__init__.py index ccf2e6e4e..e9ce14532 100644 --- a/src/packages/__init__.py +++ b/src/packages/__init__.py @@ -144,6 +144,10 @@ class ExportData: renderables: dict[RenderableType, Renderable] # The error/connection icons vbsp_conf: Property game: Game + # As objects export, they may fill this to include additional resources + # to be written to the game folder. This way it can be deferred until + # after regular resources are copied. + resources: dict[str, bytes] @attr.define @@ -604,10 +608,21 @@ def parse_package( ) return + desc: list[str] = [] + for obj in pack.info: - if obj.name in ('prerequisites', 'id', 'name', 'desc'): + if obj.name in ['prerequisites', 'id', 'name']: # Not object IDs. continue + if obj.name in ['desc', 'description']: + desc.extend(obj.as_array()) + continue + if not obj.has_children(): + LOGGER.warning( + 'Unknown package option "{}" with value "{}"!', + obj.real_name, obj.value, + ) + continue if obj.name in ('templatebrush', 'brushtemplate'): LOGGER.warning( 'TemplateBrush {} no longer needs to be defined in info.txt', @@ -661,6 +676,8 @@ def parse_package( pack.disp_name, ) + pack.desc = '\n'.join(desc) + for template in pack.fsys.walk_folder('templates'): if template.path.casefold().endswith('.vmf'): template_brush.parse_template(pack.id, template) @@ -685,7 +702,7 @@ def __init__( self.info = info self.name = name self.disp_name = disp_name - self.desc = info['desc', ''] + self.desc = '' # Filled in by parse_package. @property def enabled(self) -> bool: diff --git a/src/packages/signage.py b/src/packages/signage.py index 447e647d9..5602c6862 100644 --- a/src/packages/signage.py +++ b/src/packages/signage.py @@ -1,5 +1,7 @@ """Implements a dynamic item allowing placing the various test chamber signages.""" from __future__ import annotations + +from io import BytesIO from pathlib import Path from typing import NamedTuple, Optional, TYPE_CHECKING @@ -19,6 +21,7 @@ LOGGER = srctools.logger.get_logger(__name__) LEGEND_SIZE = (512, 1024) CELL_SIZE = 102 +SIGN_LOC = 'bee2/materials/BEE2/models/props_map_editor/signage/signage.vtf' class SignStyle(NamedTuple): @@ -215,7 +218,10 @@ def export(exp_data: ExportData) -> None: sel_icons[int(tim_id)] = sty_sign.icon exp_data.vbsp_conf.append(conf) - build_texture(exp_data.game, exp_data.selected_style, sel_icons) + exp_data.resources[SIGN_LOC] = build_texture( + exp_data.game, exp_data.selected_style, + sel_icons, + ) def _serialise(self, parent: Property, style: Style) -> Optional[SignStyle]: """Write this sign's data for the style to the provided property.""" @@ -245,7 +251,7 @@ def build_texture( game: gameMan.Game, sel_style: Style, icons: dict[int, ImgHandle], -) -> None: +) -> bytes: """Construct the legend texture for the signage.""" legend = Image.new('RGBA', LEGEND_SIZE, (0, 0, 0, 0)) @@ -285,11 +291,7 @@ def build_texture( vtf.get().copy_from(legend.tobytes(), ImageFormats.RGBA8888) vtf.clear_mipmaps() vtf.flags |= vtf.flags.ANISOTROPIC - vtf_loc = game.abs_path( - 'bee2/materials/BEE2/models/' - 'props_map_editor/signage/signage.vtf' - ) - Path(vtf_loc).parent.mkdir(parents=True, exist_ok=True) - with open(vtf_loc, 'wb') as f: - LOGGER.info('Exporting "{}"...', vtf_loc) - vtf.save(f) + with BytesIO() as buf: + vtf.save(buf) + return buf.getvalue() + diff --git a/src/packages/template_brush.py b/src/packages/template_brush.py index 234ce8e99..eef98689c 100644 --- a/src/packages/template_brush.py +++ b/src/packages/template_brush.py @@ -94,4 +94,4 @@ def write_templates(game: gameMan.Game) -> None: template_list.append(temp_el) with atomic_write(game.abs_path('bin/bee2/templates.lst'), mode='wb', overwrite=True) as f: - root.export_binary(f, fmt_name='bee_templates') + root.export_binary(f, fmt_name='bee_templates', unicode='format') diff --git a/src/precomp/conditions/__init__.py b/src/precomp/conditions/__init__.py index c98d53c3d..545a9e6a3 100644 --- a/src/precomp/conditions/__init__.py +++ b/src/precomp/conditions/__init__.py @@ -33,6 +33,7 @@ import itertools import math import random +import sys import typing import warnings from collections import defaultdict @@ -300,7 +301,25 @@ def annotation_caller( ] # For forward references and 3.7+ stringified arguments. - hints = typing.get_type_hints(func) + + # Remove 'return' temporarily so we don't parse that, since we don't care. + ann = getattr(func, '__annotations__', None) + if ann is not None: + return_val = ann.pop('return', allowed_kinds) # Sentinel + else: + return_val = None + try: + hints = typing.get_type_hints(func) + except Exception: + LOGGER.exception( + 'Could not compute type hints for function {}.{}!', + getattr(func, '__module__', ''), + func.__qualname__, + ) + sys.exit(1) # Suppress duplicate exception capture. + finally: + if ann is not None and return_val is not allowed_kinds: + ann['return'] = return_val ann_order: list[type] = [] diff --git a/src/precomp/conditions/entities.py b/src/precomp/conditions/entities.py index bcb3d7c90..a60101dd9 100644 --- a/src/precomp/conditions/entities.py +++ b/src/precomp/conditions/entities.py @@ -33,9 +33,9 @@ def res_insert_overlay(vmf: VMF, res: Property): face_str = res['face_pos', '0 0 -64'] orig_norm = Vec.from_str(res['normal', '0 0 1']) - replace_tex: dict[str, list[str]] = defaultdict(list) + replace_tex: dict[str, list[str]] = {} for prop in res.find_key('replace', []): - replace_tex[prop.name.replace('\\', '/')].append(prop.value) + replace_tex.setdefault(prop.name.replace('\\', '/'), []).append(prop.value) offset = Vec.from_str(res['offset', '0 0 0']) @@ -84,7 +84,7 @@ def insert_over(inst: Entity) -> None: random.seed('TEMP_OVERLAY_' + over['basisorigin']) mat = over['material'] try: - mat = random.choice(replace_tex[over['material'].casefold().replace('\\', '/')]) + mat = random.choice(replace_tex[mat.casefold().replace('\\', '/')]) except KeyError: pass diff --git a/src/precomp/conditions/instances.py b/src/precomp/conditions/instances.py index d1488342c..0173fc691 100644 --- a/src/precomp/conditions/instances.py +++ b/src/precomp/conditions/instances.py @@ -2,7 +2,7 @@ """ from __future__ import annotations -from typing import Union, Callable +from typing import Union, Callable, Tuple, Dict import operator import srctools.logger @@ -245,7 +245,7 @@ def res_set_inst_var(inst: Entity, res: Property): @make_result_setup('mapInstVar') -def res_map_inst_var_setup(res: Property) -> tuple[str, str, dict[str, str]]: +def res_map_inst_var_setup(res: Property) -> Tuple[str, str, Dict[str, str]]: """Pre-parse the variable table.""" table: dict[str, str] = {} res_iter = iter(res) @@ -315,7 +315,7 @@ def res_replace_instance(vmf: VMF, inst: Entity, res: Property): origin = Vec.from_str(inst['origin']) angles = Angle.from_str(inst['angles']) - if res.bool('keep_instance'): + if not res.bool('keep_instance'): inst.remove() # Do this first to free the ent ID, so the new ent has # the same one. diff --git a/src/precomp/template_brush.py b/src/precomp/template_brush.py index c45da944f..ac527e5bb 100644 --- a/src/precomp/template_brush.py +++ b/src/precomp/template_brush.py @@ -5,8 +5,7 @@ import os import random from collections import defaultdict -from collections.abc import Iterable, Iterator, Mapping -from typing import Union, Callable, Optional +from typing import Union, Callable, Optional, Tuple, Mapping, Iterable, Iterator from decimal import Decimal from enum import Enum @@ -317,8 +316,8 @@ def visgrouped( class ScalingTemplate(Mapping[ - Union[Vec, tuple[float, float, float]], - tuple[str, UVAxis, UVAxis, float] + Union[Vec, Tuple[float, float, float]], + Tuple[str, UVAxis, UVAxis, float] ]): """Represents a special version of templates, used for texturing brushes. @@ -416,7 +415,7 @@ def parse_temp_name(name) -> tuple[str, set[str]]: def load_templates(path: str) -> None: """Load in the template file, used for import_template().""" with open(path, 'rb') as f: - dmx, fmt_name, fmt_ver = DMElement.parse(f) + dmx, fmt_name, fmt_ver = DMElement.parse(f, unicode=True) if fmt_name != 'bee_templates' or fmt_ver not in [1]: raise ValueError(f'Invalid template file format "{fmt_name}" v{fmt_ver}') temp_list = dmx['temp'] diff --git a/src/pygtrie.pyi b/src/pygtrie.pyi new file mode 100644 index 000000000..b0fb9f436 --- /dev/null +++ b/src/pygtrie.pyi @@ -0,0 +1,190 @@ +"""Implements stubs for pygtrie.""" +import collections as _abc +from typing import ( + Any, Set, TypeVar, Iterator, Literal, NoReturn, + Generic, MutableMapping, overload, Mapping, Iterable, Callable, +) + + +KeyT = TypeVar('KeyT') +ValueT = TypeVar('ValueT') +T = TypeVar('T') +TrieT = TypeVar('TrieT', bound=Trie) +_EMPTY: _NoChildren + +class ShortKeyError(KeyError): ... + +class _NoChildren(Iterator[Any]): + def __bool__(self) -> Literal[False]: ... + def __nonzero__(self) -> Literal[False]: ... + def __len__(self) -> Literal[0]: ... + def __iter__(self) -> _NoChildren: ... + def iteritems(self) -> _NoChildren: ... + def sorted_items(self) -> _NoChildren: ... + def __next__(self) -> NoReturn: ... + def next(self) -> NoReturn: ... + def get(self, _step: Any) -> None: ... + def add(self, parent: Any, step: Any) -> _Node: ... + def require(self, parent: Any, step: Any) -> _Node: ... + def copy(self, _make_copy: Any, _queue: Any) -> _NoChildren: ... + def __deepcopy__(self, memo: dict) -> _NoChildren: ... + +class _OneChild(Generic[KeyT, ValueT]): + step: Any = ... + node: Any = ... + def __init__(self, step: Any, node: Any) -> None: ... + def __bool__(self) -> Literal[True]: ... + def __nonzero__(self) -> Literal[True]: ... + def __len__(self) -> Literal[1]: ... + def sorted_items(self) -> list[tuple[KeyT, ValueT]]: ... + def iteritems(self) -> Iterator[tuple[KeyT, ValueT]]: ... + def get(self, step: Any): ... + def add(self, parent: Any, step: Any): ... + def require(self, parent: Any, step: Any): ... + def delete(self, parent: Any, _step: Any) -> None: ... + def copy(self, make_copy: Any, queue: Any): ... + +class _Children(dict): + def __init__(self, *items: Any) -> None: ... + def sorted_items(self): ... + def iteritems(self): ... + def add(self, _parent: Any, step: Any): ... + def require(self, _parent: Any, step: Any): ... + def delete(self, parent: Any, step: Any) -> None: ... + def copy(self, make_copy: Any, queue: Any): ... + +class _Node: + children: Any = ... + value: Any = ... + def __init__(self) -> None: ... + def iterate(self, path: Any, shallow: Any, iteritems: Any) -> None: ... + def traverse(self, node_factory: Any, path_conv: Any, path: Any, iteritems: Any): ... + def equals(self, other: Any): ... + __bool__: Any = ... + __nonzero__: Any = ... + __hash__: Any = ... + def shallow_copy(self, make_copy: Any): ... + def copy(self, make_copy: Any): ... + +AnyNode = _Node | _NoChildren | _OneChild + +class Trie(MutableMapping[KeyT, ValueT], Generic[KeyT, ValueT]): + def __init__(self, *args: Any, **kwargs: Any) -> None: ... + def enable_sorting(self, enable: bool = ...) -> None: ... + def clear(self) -> None: ... + + @overload + def update(self: Trie[KeyT, str], m: Mapping[KeyT, ValueT], /, **kwargs: ValueT) -> None: ... + @overload + def update(self: Trie[KeyT, str], m: Iterable[tuple[KeyT, ValueT]], /, **kwargs: ValueT) -> None: ... + @overload + def update(self: Trie[KeyT, str], /, **kwargs: ValueT) -> None: ... + @overload + def update(self, m: Mapping[KeyT, ValueT], /) -> None: ... + @overload + def update(self, m: Iterable[tuple[KeyT, ValueT]], /) -> None: ... + @overload + def update(self, /) -> None: ... + + def copy(self: TrieT, __make_copy: Callable[[T], T] = ...) -> TrieT: ... + def __copy__(self: TrieT) -> TrieT: ... + def __deepcopy__(self: TrieT, memo: dict) -> TrieT: ... + + @overload + @classmethod + def fromkeys(cls: type[TrieT], keys: Iterable[KeyT]) -> TrieT[KeyT, None]: ... + @overload + @classmethod + def fromkeys(cls: type[TrieT], keys: Iterable[KeyT], value: ValueT) -> TrieT[KeyT, ValueT]: ... + + def __iter__(self) -> list[KeyT]: ... + def iteritems(self, prefix: KeyT = ..., shallow: bool = ...) -> Iterator[tuple[KeyT, ValueT]]: ... + def iterkeys(self, prefix: KeyT = ..., shallow: bool = ...) -> Iterator[KeyT]: ... + def itervalues(self, prefix: KeyT = ..., shallow: bool = ...) -> Iterator[ValueT]: ... + def items(self, prefix: KeyT = ..., shallow: bool = ...) -> list[tuple[KeyT, ValueT]]: ... + def keys(self, prefix: KeyT = ..., shallow: bool = ...) -> list[KeyT]: ... + def values(self, prefix: KeyT = ..., shallow: bool = ...) -> list[ValueT]: ... + def __len__(self) -> int: ... + def __bool__(self) -> bool: ... + def __nonzero__(self) -> bool: ... + __hash__: None + HAS_VALUE: int + HAS_SUBTRIE: int + def has_node(self, key: KeyT): ... + def has_key(self, key: KeyT): ... + def has_subtrie(self, key: KeyT): ... + # TODO: slice can't specify it must always be slice(KeyT, None, None) + def __getitem__(self, key_or_slice: KeyT | slice): ... + def __setitem__(self, key_or_slice: KeyT | slice, value: ValueT) -> None: ... + def setdefault(self, key: KeyT, default: ValueT = None) -> ValueT: ... + @overload + def pop(self, key: KeyT) -> ValueT: ... + @overload + def pop(self, key: KeyT, default: ValueT | T = ...) -> ValueT | T: ... + def popitem(self) -> tuple[KeyT, ValueT]: ... + def __delitem__(self, key_or_slice: KeyT | slice) -> None: ... + + class _NoneStep: + def __bool__(self) -> Literal[False]: ... + def __nonzero__(self) -> Literal[False]: ... + def get(self, default: T = None) -> T: ... + is_set: Literal[False] + has_subtrie: Literal[False] + def key(self) -> None: ... + def value(self) -> None: ... + def __getitem__(self, index: int) -> None: ... + + class _Step(_NoneStep, Generic[KeyT, ValueT]): + def __init__(self, trie: Trie, path: KeyT, pos: int, node: AnyNode) -> None: ... + def __bool__(self) -> bool: ... + def __nonzero__(self) -> bool: ... + @property + def is_set(self) -> bool: ... + @property + def has_subtrie(self) -> bool: ... + def get(self, default: T = None) -> ValueT | T: ... + def set(self, value: ValueT) -> None: ... + def setdefault(self, value: ValueT) -> ValueT: ... + @property + def key(self) -> KeyT: ... + @property + def value(self) -> ValueT: ... + @value.setter + def value(self, value: ValueT) -> None: ... + + def walk_towards(self, key: KeyT) -> Iterator[_Step[KeyT, ValueT]]: ... + def prefixes(self, key: KeyT) -> Iterator[_Step[KeyT, ValueT]]: ... + def shortest_prefix(self, key: KeyT): ... + def longest_prefix(self, key: KeyT): ... + def __eq__(self, other: Trie[KeyT, ValueT]) -> bool: ... + def __ne__(self, other: Trie[KeyT, ValueT]) -> bool: ... + def traverse(self, node_factory: Callable[..., T], prefix: KeyT = ...) -> T: ... + +class CharTrie(Trie[str, ValueT], Generic[ValueT]): ... + +class StringTrie(Trie[str, ValueT], Generic[ValueT]): + def __init__(self, *args: Any, separator: str='/', **kwargs: Any) -> None: ... + + @overload + @classmethod + def fromkeys(cls, keys: Iterable[str], *, separator: str = ...) -> StringTrie[None]: ... + @overload + @classmethod + def fromkeys(cls, keys: Iterable[str], value: ValueT, separator: str = ...) -> StringTrie[ValueT]: ... + +class PrefixSet(Set[KeyT], Generic[KeyT]): + # TODO: Used as factory(**kwargs), but can't express that. + def __init__(self, iterable: Iterable[KeyT] = ..., factory: Callable[..., Trie] = ..., **kwargs: Any) -> None: ... + def copy(self) -> PrefixSet[KeyT]: ... + def __copy__(self) -> PrefixSet[KeyT]: ... + def __deepcopy__(self, memo: dict) -> PrefixSet[KeyT]: ... + def clear(self) -> None: ... + def __contains__(self, key: KeyT) -> bool: ... + def __iter__(self) -> Iterator[KeyT]: ... + def iter(self, prefix: KeyT = ...) -> Iterator[KeyT]: ... + def __len__(self) -> int: ... + def add(self, value: KeyT) -> None: ... + # Not implemented. + def discard(self, value: KeyT) -> NoReturn: ... + def remove(self, value: KeyT) -> NoReturn: ... + def pop(self) -> NoReturn: ... diff --git a/src/utils.py b/src/utils.py index 16f531f4e..8b4170b43 100644 --- a/src/utils.py +++ b/src/utils.py @@ -375,8 +375,10 @@ def __init__(self, pack_id: str, path: str) -> None: self.path = path.replace('\\', '/') @classmethod - def parse(cls, uri: str, def_package: str) -> PackagePath: + def parse(cls, uri: str | PackagePath, def_package: str) -> PackagePath: """Parse a string into a path. If a package isn't provided, the default is used.""" + if isinstance(uri, PackagePath): + return uri if ':' in uri: return cls(*uri.split(':', 1)) else: