17
17
PATTERN = r"""
18
18
(
19
19
(?<!pd\.)(?<!\w) # check class_name start with pd. or character
20
- {class_name} \( # match DataFrame but not pd.DataFrame or tm.makeDataFrame
20
+ ([A-Z]\w+) \( # match DataFrame but not pd.DataFrame or tm.makeDataFrame
21
21
.* # match anything
22
- pd\.{class_name} \( # only match e.g. pd.DataFrame
22
+ pd\.\2 \( # only match e.g. pd.DataFrame
23
23
)|
24
24
(
25
- pd\.{class_name} \( # only match e.g. pd.DataFrame
25
+ pd\.([A-Z]\w+) \( # only match e.g. pd.DataFrame
26
26
.* # match anything
27
27
(?<!pd\.)(?<!\w) # check class_name start with pd. or character
28
- {class_name} \( # match DataFrame but not pd.DataFrame or tm.makeDataFrame
28
+ \4 \( # match DataFrame but not pd.DataFrame or tm.makeDataFrame
29
29
)
30
30
"""
31
- CLASS_NAMES = (
32
- "Series" ,
33
- "DataFrame" ,
34
- "Index" ,
35
- "MultiIndex" ,
36
- "Timestamp" ,
37
- "Timedelta" ,
38
- "TimedeltaIndex" ,
39
- "DatetimeIndex" ,
40
- "Categorical" ,
41
- )
42
31
ERROR_MESSAGE = "Found both `pd.{class_name}` and `{class_name}` in {path}"
43
32
44
33
@@ -47,16 +36,22 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
47
36
parser .add_argument ("paths" , nargs = "*" , type = Path )
48
37
args = parser .parse_args (argv )
49
38
50
- for class_name in CLASS_NAMES :
51
- pattern = re .compile (
52
- PATTERN .format (class_name = class_name ).encode (),
53
- flags = re .MULTILINE | re .DOTALL | re .VERBOSE ,
54
- )
55
- for path in args .paths :
56
- contents = path .read_bytes ()
57
- match = pattern .search (contents )
58
- assert match is None , ERROR_MESSAGE .format (
59
- class_name = class_name , path = str (path )
39
+ pattern = re .compile (
40
+ PATTERN .encode (),
41
+ flags = re .MULTILINE | re .DOTALL | re .VERBOSE ,
42
+ )
43
+ for path in args .paths :
44
+ contents = path .read_bytes ()
45
+ match = pattern .search (contents )
46
+ if match is None :
47
+ continue
48
+ if match .group (2 ) is not None :
49
+ raise AssertionError (
50
+ ERROR_MESSAGE .format (class_name = match .group (2 ).decode (), path = str (path ))
51
+ )
52
+ if match .group (4 ) is not None :
53
+ raise AssertionError (
54
+ ERROR_MESSAGE .format (class_name = match .group (4 ).decode (), path = str (path ))
60
55
)
61
56
62
57
0 commit comments