Skip to content

Commit 09830d2

Browse files
SchamperSchamper
authored
Fix enum and flag repr in struct reprs (#101)
1 parent 046dfa5 commit 09830d2

File tree

5 files changed

+138
-43
lines changed

5 files changed

+138
-43
lines changed

‎dissect/cstruct/types/structure.py

+10-4
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import io
44
from contextlib import contextmanager
5+
from enum import Enum
56
from functools import lru_cache
67
from operator import attrgetter
78
from textwrap import dedent
@@ -380,10 +381,15 @@ def __getitem__(self, item: str) -> Any:
380381
return getattr(self, item)
381382

382383
def __repr__(self) -> str:
383-
values = [
384-
f"{k}={hex(self[k]) if (issubclass(f.type, int) and not issubclass(f.type, Pointer)) else repr(self[k])}"
385-
for k, f in self.__class__.fields.items()
386-
]
384+
values = []
385+
for name, field in self.__class__.fields.items():
386+
value = self[name]
387+
if issubclass(field.type, int) and not issubclass(field.type, (Pointer, Enum)):
388+
value = hex(value)
389+
else:
390+
value = repr(value)
391+
values.append(f"{name}={value}")
392+
387393
return f"<{self.__class__.__name__} {' '.join(values)}>"
388394

389395

‎dissect/cstruct/utils.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@
33
import pprint
44
import string
55
import sys
6+
from enum import Enum
67
from typing import Iterator
78

9+
from dissect.cstruct.types.pointer import Pointer
810
from dissect.cstruct.types.structure import Structure
911

1012
COLOR_RED = "\033[1;31m"
@@ -158,7 +160,7 @@ def _dumpstruct(
158160
ci += 1
159161

160162
value = getattr(structure, field.name)
161-
if isinstance(value, str):
163+
if isinstance(value, (str, Pointer, Enum)):
162164
value = repr(value)
163165
elif isinstance(value, int):
164166
value = hex(value)

‎tests/test_types_enum.py

+63-23
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,66 @@ def test_enum_eof(TestEnum: type[Enum]) -> None:
9898
TestEnum[None](b"\x01")
9999

100100

101+
def test_enum_same_value_different_type(cs: cstruct, compiled: bool) -> None:
102+
cdef = """
103+
enum Test16 : uint16 {
104+
A = 0x1,
105+
B = 0x2
106+
};
107+
108+
enum Test24 : uint24 {
109+
A = 0x1,
110+
B = 0x2
111+
};
112+
113+
enum Test32 : uint32 {
114+
A = 0x1,
115+
B = 0x2
116+
};
117+
"""
118+
cs.load(cdef, compiled=compiled)
119+
120+
obj = {
121+
cs.Test16.A: "Test16.A",
122+
cs.Test16.B: "Test16.B",
123+
cs.Test24.A: "Test24.A",
124+
cs.Test24.B: "Test24.B",
125+
}
126+
127+
assert obj[cs.Test16.A] == "Test16.A"
128+
assert obj[cs.Test16(2)] == "Test16.B"
129+
assert obj[cs.Test24(1)] == "Test24.A"
130+
assert obj[cs.Test24.B] == "Test24.B"
131+
132+
with pytest.raises(KeyError):
133+
obj[cs.Test32.A]
134+
135+
136+
def test_enum_str_repr(TestEnum: type[Enum]) -> None:
137+
assert repr(TestEnum.A) == "<Test.A: 1>"
138+
assert str(TestEnum.A) == "Test.A"
139+
assert repr(TestEnum(69)) == "<Test: 69>"
140+
assert str(TestEnum(69)) == "Test.69"
141+
142+
143+
def test_enum_str_repr_in_struct(cs: cstruct, compiled: bool) -> None:
144+
cdef = """
145+
enum Test16 : uint16 {
146+
A = 0x1,
147+
B = 0x2
148+
};
149+
150+
struct test {
151+
Test16 a;
152+
};
153+
"""
154+
cs.load(cdef, compiled=compiled)
155+
156+
obj = cs.test(b"\x02\x00")
157+
assert repr(obj) == "<test a=<Test16.B: 2>>"
158+
assert str(obj) == "<test a=<Test16.B: 2>>"
159+
160+
101161
def test_enum_struct(cs: cstruct, compiled: bool) -> None:
102162
cdef = """
103163
enum Test16 : uint16 {
@@ -174,26 +234,6 @@ def test_enum_struct(cs: cstruct, compiled: bool) -> None:
174234
assert cs.test_expr(buf).expr == [cs.Test16.A, cs.Test16.B]
175235
assert cs.test_expr(size=1, expr=[cs.Test16.A, cs.Test16.B]).dumps() == buf
176236

177-
obj = {
178-
cs.Test16.A: "Test16.A",
179-
cs.Test16.B: "Test16.B",
180-
cs.Test24.A: "Test24.A",
181-
cs.Test24.B: "Test24.B",
182-
}
183-
184-
assert obj[cs.Test16.A] == "Test16.A"
185-
assert obj[cs.Test16(2)] == "Test16.B"
186-
assert obj[cs.Test24(1)] == "Test24.A"
187-
assert obj[cs.Test24.B] == "Test24.B"
188-
189-
with pytest.raises(KeyError):
190-
obj[cs.Test32.A]
191-
192-
assert repr(cs.Test16.A) == "<Test16.A: 1>"
193-
assert str(cs.Test16.A) == "Test16.A"
194-
assert repr(cs.Test16(69)) == "<Test16: 69>"
195-
assert str(cs.Test16(69)) == "Test16.69"
196-
197237

198238
def test_enum_comments(cs: cstruct) -> None:
199239
cdef = """
@@ -259,15 +299,15 @@ def test_enum_name(cs: cstruct, compiled: bool) -> None:
259299
Color = cs.Color
260300
Pixel = cs.Pixel
261301

262-
pixel = Pixel(b"\xFF\x0A\x01\x00\xAA\xBB\xCC\xDD")
302+
pixel = Pixel(b"\xff\x0a\x01\x00\xaa\xbb\xcc\xdd")
263303
assert pixel.x == 255
264304
assert pixel.y == 10
265305
assert pixel.color.name == "RED"
266306
assert pixel.color.value == Color.RED
267307
assert pixel.color.value == 1
268308
assert pixel.hue == 0xDDCCBBAA
269309

270-
pixel = Pixel(b"\x00\x00\xFF\x00\xAA\xBB\xCC\xDD")
310+
pixel = Pixel(b"\x00\x00\xff\x00\xaa\xbb\xcc\xdd")
271311
assert pixel.color.name is None
272312
assert pixel.color.value == 0xFF
273313
assert repr(pixel.color) == "<Color: 255>"
@@ -350,7 +390,7 @@ def test_enum_anonymous_struct(cs: cstruct, compiled: bool) -> None:
350390

351391
test = cs.test
352392

353-
t = test(b"\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0A\x00\x00\x00")
393+
t = test(b"\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0a\x00\x00\x00")
354394
assert t.arr == [255, 0, 0, 10]
355395

356396

‎tests/test_types_flag.py

+37-15
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,42 @@ def test_flag_operator(TestFlag: type[Flag]) -> None:
6262
assert TestFlag[2]([TestFlag.B, (TestFlag.A | TestFlag.B)]).dumps() == b"\x02\x03"
6363

6464

65+
def test_flag_str_repr(TestFlag: type[Flag]) -> None:
66+
if PY_311:
67+
assert repr(TestFlag.A | TestFlag.B) == "<Test.A|B: 3>"
68+
assert str(TestFlag.A | TestFlag.B) == "Test.A|B"
69+
assert repr(TestFlag(69)) == "<Test.A|68: 69>"
70+
assert str(TestFlag(69)) == "Test.A|68"
71+
else:
72+
assert repr(TestFlag.A | TestFlag.B) == "<Test.B|A: 3>"
73+
assert str(TestFlag.A | TestFlag.B) == "Test.B|A"
74+
assert repr(TestFlag(69)) == "<Test.64|4|A: 69>"
75+
assert str(TestFlag(69)) == "Test.64|4|A"
76+
77+
78+
def test_flag_str_repr_in_struct(cs: cstruct, compiled: bool) -> None:
79+
cdef = """
80+
flag Test : uint16 {
81+
A,
82+
B
83+
};
84+
85+
struct test {
86+
Test a;
87+
};
88+
"""
89+
cs.load(cdef, compiled=compiled)
90+
91+
obj = cs.test(b"\x03\x00")
92+
93+
if PY_311:
94+
assert repr(obj) == "<test a=<Test.A|B: 3>>"
95+
assert str(obj) == "<test a=<Test.A|B: 3>>"
96+
else:
97+
assert repr(obj) == "<test a=<Test.B|A: 3>>"
98+
assert str(obj) == "<test a=<Test.B|A: 3>>"
99+
100+
65101
def test_flag_struct(cs: cstruct) -> None:
66102
cdef = """
67103
flag Test {
@@ -101,20 +137,6 @@ def test_flag_struct(cs: cstruct) -> None:
101137
assert bool(cs.Test(1)) is True
102138

103139
assert cs.Test.a | cs.Test.b == 3
104-
if PY_311:
105-
assert repr(cs.Test.c | cs.Test.d) == "<Test.c|d: 12>"
106-
assert str(cs.Test.c | cs.Test.d) == "Test.c|d"
107-
assert repr(cs.Test.a | cs.Test.b) == "<Test.a|b: 3>"
108-
assert str(cs.Test.a | cs.Test.b) == "Test.a|b"
109-
assert repr(cs.Test(69)) == "<Test.a|c|64: 69>"
110-
assert str(cs.Test(69)) == "Test.a|c|64"
111-
else:
112-
assert repr(cs.Test.c | cs.Test.d) == "<Test.d|c: 12>"
113-
assert str(cs.Test.c | cs.Test.d) == "Test.d|c"
114-
assert repr(cs.Test.a | cs.Test.b) == "<Test.b|a: 3>"
115-
assert str(cs.Test.a | cs.Test.b) == "Test.b|a"
116-
assert repr(cs.Test(69)) == "<Test.64|c|a: 69>"
117-
assert str(cs.Test(69)) == "Test.64|c|a"
118140
assert cs.Test(2) == cs.Test.b
119141
assert cs.Test(3) == cs.Test.a | cs.Test.b
120142
assert cs.Test.c & 12 == cs.Test.c
@@ -231,5 +253,5 @@ def test_flag_anonymous_struct(cs: cstruct, compiled: bool) -> None:
231253

232254
test = cs.test
233255

234-
t = test(b"\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0A\x00\x00\x00")
256+
t = test(b"\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0a\x00\x00\x00")
235257
assert t.arr == [255, 0, 0, 10]

‎tests/test_utils.py

+25
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,31 @@ def test_dumpstruct_anonymous(cs: cstruct, capsys: pytest.CaptureFixture, compil
8686
assert str(excinfo.value) == "Invalid output argument: 'generator' (should be 'print' or 'string')."
8787

8888

89+
def test_dumpstruct_enum(cs: cstruct, compiled: bool) -> None:
90+
cdef = """
91+
enum Test16 : uint16 {
92+
A = 0x1,
93+
B = 0x2
94+
};
95+
96+
struct test {
97+
Test16 testval;
98+
};
99+
"""
100+
cs.load(cdef, compiled=compiled)
101+
102+
assert verify_compiled(cs.test, compiled)
103+
104+
buf = b"\x02\x00"
105+
obj = cs.test(buf)
106+
107+
out1 = utils.dumpstruct(cs.test, buf, output="string")
108+
out2 = utils.dumpstruct(obj, output="string")
109+
110+
assert "<Test16.B: 2>" in out1
111+
assert "<Test16.B: 2>" in out2
112+
113+
89114
def test_pack_unpack() -> None:
90115
endian = "little"
91116
sign = False

0 commit comments

Comments
 (0)