aws_fully_qualified_names/languages/python/
mod.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4mod boto3_service_name_to_module_name;
5mod builtins;
6mod traverse;
7
8use super::*;
9use crate::languages::ast_visitor::{DeclUseVisitor, RelevantName, ScopeVisitor};
10use traverse::traverse_ast;
11#[cfg(target_arch = "wasm32")]
12use wasm_bindgen::prelude::*;
13
14#[cfg_attr(target_arch = "wasm32", wasm_bindgen)]
15pub struct Python;
16
17impl TreeSitterLanguage for Python {
18    fn find_symbol_declarations_and_usages<'code>(
19        #[cfg(not(target_arch = "wasm32"))] cursor: &mut TreeCursor<'code>,
20        #[cfg(target_arch = "wasm32")] cursor: &mut TreeCursor,
21        code: &'code [u8],
22    ) -> ScopedDeclarationsAndUsages {
23        traverse_ast(cursor, code, DeclUseVisitor::new())
24    }
25
26    #[cfg(not(target_arch = "wasm32"))]
27    fn grammar() -> Grammar {
28        tree_sitter_python::language()
29    }
30}
31
32#[cfg_attr(target_arch = "wasm32", wasm_bindgen)]
33impl Python {
34    #[cfg(target_arch = "wasm32")]
35    #[wasm_bindgen(js_name = "findNames")]
36    pub async fn find_names(code: String) -> JsValue {
37        Parser::init().await;
38        let grammar =
39            crate::wasm::tree_sitter::load(include_bytes!(std::env!("PYTHON_LANGUAGE"))).await;
40        let parser = Parser::new();
41        parser.set_language(grammar);
42        let ast = parser.parse(&code);
43        let names = <Self as Language>::find_names(&ast, &code);
44        serde_wasm_bindgen::to_value(&names).unwrap()
45    }
46
47    #[cfg(target_arch = "wasm32")]
48    #[wasm_bindgen(js_name = "findNamesWithInExtent")]
49    pub async fn find_names_within_extent(code: String, extent: Extent) -> JsValue {
50        Parser::init().await;
51        let grammar =
52            crate::wasm::tree_sitter::load(include_bytes!(std::env!("PYTHON_LANGUAGE"))).await;
53        let parser = Parser::new();
54        parser.set_language(grammar);
55        let ast = parser.parse(&code);
56        let names = <Self as Language>::find_names(&ast, &code);
57        let extent_names = names.within_extent(&extent);
58        serde_wasm_bindgen::to_value(&extent_names).unwrap()
59    }
60
61    #[cfg(not(target_arch = "wasm32"))]
62    pub fn find_relevant_scope(
63        code: &str,
64        relevant_names: Vec<RelevantName>,
65    ) -> Result<RelevantScope> {
66        let code = code.as_bytes();
67        let ast = Python::parse(code)?;
68        let mut cursor = ast.walk();
69        let node = cursor.node();
70        let byte_range = node.byte_range();
71        let extent = node.into();
72        let (relevant_scope, snippet_extent, relevant_names_occurrences) = traverse_ast(
73            &mut cursor,
74            code,
75            ScopeVisitor::new(
76                Scope {
77                    extent,
78                    byte_range,
79                    name: None,
80                },
81                relevant_names,
82            ),
83        );
84        Ok(RelevantScope {
85            snippet: String::from_utf8_lossy(&code[relevant_scope]).to_string(),
86            snippet_extent,
87            relevant_names_occurrences,
88        })
89    }
90}
91
92#[cfg(all(test, not(target_arch = "wasm32")))]
93mod test {
94    use super::*;
95    use crate::languages::test_macros::*;
96
97    #[test]
98    fn finds_unicode_name() {
99        let code = r#"# a function that prints
100def foo_by̆r():
101    # 😊 smiley face
102    print("invoked foo_by̆r")
103
104foo_by̆r()
105print("done")"#;
106
107        let names = Python::find_names(code).expect("Cannot find names in test code");
108        print!("{:#?}", names);
109        assert_eq!(
110            names,
111            Names {
112                fully_qualified: DeclUse {
113                    declared_symbols: vec![],
114                    used_symbols: vec![],
115                },
116                simple: DeclUse {
117                    declared_symbols: vec![SimpleName {
118                        symbol: "foo_by\u{306}r".to_string(),
119                        extent: Some(Extent {
120                            start: Location {
121                                line: 1,
122                                character: 4,
123                            },
124                            end: Location {
125                                line: 1,
126                                character: 13,
127                            },
128                        },),
129                    },],
130                    used_symbols: vec![
131                        SimpleName {
132                            symbol: "print".to_string(),
133                            extent: Some(Extent {
134                                start: Location {
135                                    line: 3,
136                                    character: 4,
137                                },
138                                end: Location {
139                                    line: 3,
140                                    character: 9,
141                                },
142                            },),
143                        },
144                        SimpleName {
145                            symbol: "foo_by\u{306}r".to_string(),
146                            extent: Some(Extent {
147                                start: Location {
148                                    line: 5,
149                                    character: 0,
150                                },
151                                end: Location {
152                                    line: 5,
153                                    character: 9,
154                                },
155                            },),
156                        },
157                        SimpleName {
158                            symbol: "print".to_string(),
159                            extent: Some(Extent {
160                                start: Location {
161                                    line: 6,
162                                    character: 0,
163                                },
164                                end: Location {
165                                    line: 6,
166                                    character: 5,
167                                },
168                            },),
169                        },
170                    ],
171                },
172                external_simple: DeclUse {
173                    declared_symbols: vec![],
174                    used_symbols: vec![],
175                },
176            }
177        );
178    }
179
180    #[test]
181    fn finds_names_in_basic_constructs() {
182        let code = r#"from difflib import Differ
183hello = "HelloWorld"
184print(hello)
185left, right = ("left" + 45), "right" + 45.0
186
187def diff(a, b="right"):
188    import re as hello
189    pattern = hello.compile('[^4]+')
190    pattern.sub('5', b)
191    return Differ().compare(a=a, b=b)
192
193diff(left, "left")
194
195import diff
196
197raise diff.diff(left, "left")
198if 3 > 9:
199    pass
200elif True:
201    pass
202else:
203    ...
204"#;
205        let names = Python::find_names(code).expect("Cannot find names in test code");
206
207        assert_eq!(
208            names,
209            Names {
210                fully_qualified: DeclUse {
211                    declared_symbols: vec![
212                        FullyQualifiedName {
213                            source: sym!(difflib),
214                            symbol: sym!(Differ),
215                            extent: extent!(0:20-0:26),
216                        },
217                        FullyQualifiedName {
218                            source: sym!(re),
219                            symbol: vec![],
220                            extent: extent!(6:11-6:13),
221                        },
222                        FullyQualifiedName {
223                            source: sym!(diff),
224                            symbol: vec![],
225                            extent: extent!(13:7-13:11),
226                        }
227                    ],
228                    used_symbols: vec![
229                        FullyQualifiedName {
230                            source: sym!(re),
231                            symbol: vec![],
232                            extent: extent!(7:14-7:19),
233                        },
234                        FullyQualifiedName {
235                            source: sym!(re),
236                            symbol: sym!(compile),
237                            extent: extent!(7:20-7:27),
238                        },
239                        FullyQualifiedName {
240                            source: sym!(difflib),
241                            symbol: sym!(Differ),
242                            extent: extent!(9:11-9:17),
243                        },
244                        FullyQualifiedName {
245                            source: sym!(difflib),
246                            symbol: sym!(Differ.compare),
247                            extent: extent!(9:20-9:27),
248                        },
249                        FullyQualifiedName {
250                            source: sym!(diff),
251                            symbol: vec![],
252                            extent: extent!(15:6-15:10),
253                        },
254                        FullyQualifiedName {
255                            source: sym!(diff),
256                            symbol: sym!(diff),
257                            extent: extent!(15:11-15:15),
258                        }
259                    ]
260                },
261                simple: DeclUse {
262                    declared_symbols: vec![
263                        SimpleName {
264                            symbol: "difflib".to_string(),
265                            extent: optional_extent!(0:5-0:12),
266                        },
267                        SimpleName {
268                            symbol: "Differ".to_string(),
269                            extent: optional_extent!(0:20-0:26),
270                        },
271                        SimpleName {
272                            symbol: "hello".to_string(),
273                            extent: optional_extent!(1:0-1:5),
274                        },
275                        SimpleName {
276                            symbol: "left".to_string(),
277                            extent: optional_extent!(3:0-3:4),
278                        },
279                        SimpleName {
280                            symbol: "right".to_string(),
281                            extent: optional_extent!(3:6-3:11),
282                        },
283                        SimpleName {
284                            symbol: "diff".to_string(),
285                            extent: optional_extent!(5:4-5:8),
286                        },
287                        SimpleName {
288                            symbol: "a".to_string(),
289                            extent: optional_extent!(5:9-5:10),
290                        },
291                        SimpleName {
292                            symbol: "b".to_string(),
293                            extent: optional_extent!(5:12-5:13),
294                        },
295                        SimpleName {
296                            symbol: "re".to_string(),
297                            extent: optional_extent!(6:11-6:13),
298                        },
299                        SimpleName {
300                            symbol: "pattern".to_string(),
301                            extent: optional_extent!(7:4-7:11),
302                        },
303                        SimpleName {
304                            symbol: "diff".to_string(),
305                            extent: optional_extent!(13:7-13:11),
306                        }
307                    ],
308                    used_symbols: vec![
309                        SimpleName {
310                            symbol: "print".to_string(),
311                            extent: optional_extent!(2:0-2:5),
312                        },
313                        SimpleName {
314                            symbol: "hello".to_string(),
315                            extent: optional_extent!(2:6-2:11),
316                        },
317                        SimpleName {
318                            symbol: "re".to_string(),
319                            extent: optional_extent!(7:14-7:19),
320                        },
321                        SimpleName {
322                            symbol: "re".to_string(),
323                            extent: None
324                        },
325                        SimpleName {
326                            symbol: "compile".to_string(),
327                            extent: optional_extent!(7:20-7:27),
328                        },
329                        SimpleName {
330                            symbol: "pattern".to_string(),
331                            extent: optional_extent!(8:4-8:11),
332                        },
333                        SimpleName {
334                            symbol: "sub".to_string(),
335                            extent: optional_extent!(8:12-8:15),
336                        },
337                        SimpleName {
338                            symbol: "b".to_string(),
339                            extent: optional_extent!(8:21-8:22),
340                        },
341                        SimpleName {
342                            symbol: "difflib".to_string(),
343                            extent: None
344                        },
345                        SimpleName {
346                            symbol: "Differ".to_string(),
347                            extent: optional_extent!(9:11-9:17),
348                        },
349                        SimpleName {
350                            symbol: "difflib".to_string(),
351                            extent: None,
352                        },
353                        SimpleName {
354                            symbol: "Differ".to_string(),
355                            extent: None,
356                        },
357                        SimpleName {
358                            symbol: "compare".to_string(),
359                            extent: optional_extent!(9:20-9:27),
360                        },
361                        SimpleName {
362                            symbol: "a".to_string(),
363                            extent: optional_extent!(9:30-9:31),
364                        },
365                        SimpleName {
366                            symbol: "b".to_string(),
367                            extent: optional_extent!(9:35-9:36),
368                        },
369                        SimpleName {
370                            symbol: "diff".to_string(),
371                            extent: optional_extent!(11:0-11:4),
372                        },
373                        SimpleName {
374                            symbol: "left".to_string(),
375                            extent: optional_extent!(11:5-11:9),
376                        },
377                        SimpleName {
378                            symbol: "diff".to_string(),
379                            extent: optional_extent!(15:6-15:10),
380                        },
381                        SimpleName {
382                            symbol: "diff".to_string(),
383                            extent: None
384                        },
385                        SimpleName {
386                            symbol: "diff".to_string(),
387                            extent: optional_extent!(15:11-15:15),
388                        },
389                        SimpleName {
390                            symbol: "left".to_string(),
391                            extent: optional_extent!(15:16-15:20),
392                        }
393                    ]
394                },
395                external_simple: DeclUse {
396                    declared_symbols: vec![
397                        SimpleName {
398                            symbol: "difflib".to_string(),
399                            extent: optional_extent!(0:5-0:12),
400                        },
401                        SimpleName {
402                            symbol: "Differ".to_string(),
403                            extent: optional_extent!(0:20-0:26),
404                        },
405                        SimpleName {
406                            symbol: "re".to_string(),
407                            extent: optional_extent!(6:11-6:13),
408                        },
409                        SimpleName {
410                            symbol: "diff".to_string(),
411                            extent: optional_extent!(13:7-13:11),
412                        }
413                    ],
414                    used_symbols: vec![
415                        SimpleName {
416                            symbol: "re".to_string(),
417                            extent: optional_extent!(7:14-7:19),
418                        },
419                        SimpleName {
420                            symbol: "re".to_string(),
421                            extent: None
422                        },
423                        SimpleName {
424                            symbol: "compile".to_string(),
425                            extent: optional_extent!(7:20-7:27),
426                        },
427                        SimpleName {
428                            symbol: "sub".to_string(),
429                            extent: optional_extent!(8:12-8:15),
430                        },
431                        SimpleName {
432                            symbol: "difflib".to_string(),
433                            extent: None
434                        },
435                        SimpleName {
436                            symbol: "Differ".to_string(),
437                            extent: optional_extent!(9:11-9:17),
438                        },
439                        SimpleName {
440                            symbol: "difflib".to_string(),
441                            extent: None,
442                        },
443                        SimpleName {
444                            symbol: "Differ".to_string(),
445                            extent: None,
446                        },
447                        SimpleName {
448                            symbol: "compare".to_string(),
449                            extent: optional_extent!(9:20-9:27),
450                        },
451                        SimpleName {
452                            symbol: "diff".to_string(),
453                            extent: optional_extent!(15:6-15:10),
454                        },
455                        SimpleName {
456                            symbol: "diff".to_string(),
457                            extent: None
458                        },
459                        SimpleName {
460                            symbol: "diff".to_string(),
461                            extent: optional_extent!(15:11-15:15),
462                        }
463                    ]
464                }
465            }
466        );
467    }
468
469    #[test]
470    fn finds_names_in_nested_comprehension_expressions() {
471        let code = r#"from xlib import x
472[x for x in range(x) for x in range(x) if x > 3]
473"#;
474        let names = Python::find_names(code).expect("Cannot find names in test code");
475
476        assert_eq!(
477            names,
478            Names {
479                fully_qualified: DeclUse {
480                    declared_symbols: vec![FullyQualifiedName {
481                        source: sym!(xlib),
482                        symbol: sym!(x),
483                        extent: extent!(0:17-0:18),
484                    }],
485                    used_symbols: vec![FullyQualifiedName {
486                        source: sym!(xlib),
487                        symbol: sym!(x),
488                        extent: extent!(1:18-1:19),
489                    }]
490                },
491                simple: DeclUse {
492                    declared_symbols: vec![
493                        SimpleName {
494                            symbol: "xlib".to_string(),
495                            extent: optional_extent!(0:5-0:9),
496                        },
497                        SimpleName {
498                            symbol: "x".to_string(),
499                            extent: optional_extent!(0:17-0:18),
500                        },
501                        SimpleName {
502                            symbol: "x".to_string(),
503                            extent: optional_extent!(1:7-1:8),
504                        },
505                        SimpleName {
506                            symbol: "x".to_string(),
507                            extent: optional_extent!(1:25-1:26),
508                        }
509                    ],
510                    used_symbols: vec![
511                        SimpleName {
512                            symbol: "range".to_string(),
513                            extent: optional_extent!(1:12-1:17),
514                        },
515                        SimpleName {
516                            symbol: "xlib".to_string(),
517                            extent: optional_extent!(),
518                        },
519                        SimpleName {
520                            symbol: "x".to_string(),
521                            extent: optional_extent!(1:18-1:19),
522                        },
523                        SimpleName {
524                            symbol: "range".to_string(),
525                            extent: optional_extent!(1:30-1:35),
526                        },
527                        SimpleName {
528                            symbol: "x".to_string(),
529                            extent: optional_extent!(1:36-1:37),
530                        },
531                        SimpleName {
532                            symbol: "x".to_string(),
533                            extent: optional_extent!(1:42-1:43),
534                        },
535                        SimpleName {
536                            symbol: "x".to_string(),
537                            extent: optional_extent!(1:1-1:2),
538                        }
539                    ]
540                },
541                external_simple: DeclUse {
542                    declared_symbols: vec![
543                        SimpleName {
544                            symbol: "xlib".to_string(),
545                            extent: optional_extent!(0:5-0:9),
546                        },
547                        SimpleName {
548                            symbol: "x".to_string(),
549                            extent: optional_extent!(0:17-0:18),
550                        }
551                    ],
552                    used_symbols: vec![
553                        SimpleName {
554                            symbol: "xlib".to_string(),
555                            extent: optional_extent!(),
556                        },
557                        SimpleName {
558                            symbol: "x".to_string(),
559                            extent: optional_extent!(1:18-1:19),
560                        }
561                    ]
562                },
563            }
564        );
565    }
566
567    #[test]
568    fn finds_names_through_global_statements() {
569        let code = r#"a = 4
570def fn():
571    a = 6
572    def prt():
573        global a
574        print(a)
575
576"#;
577        let names = Python::find_names(code).expect("Cannot find names in test code");
578
579        assert_eq!(
580            names,
581            Names {
582                fully_qualified: DeclUse {
583                    declared_symbols: vec![],
584                    used_symbols: vec![]
585                },
586                simple: DeclUse {
587                    declared_symbols: vec![
588                        SimpleName {
589                            symbol: "a".to_string(),
590                            extent: optional_extent!(0:0-0:1),
591                        },
592                        SimpleName {
593                            symbol: "fn".to_string(),
594                            extent: optional_extent!(1:4-1:6),
595                        },
596                        SimpleName {
597                            symbol: "a".to_string(),
598                            extent: optional_extent!(2:4-2:5),
599                        },
600                        SimpleName {
601                            symbol: "prt".to_string(),
602                            extent: optional_extent!(3:8-3:11),
603                        }
604                    ],
605                    used_symbols: vec![
606                        SimpleName {
607                            symbol: "print".to_string(),
608                            extent: optional_extent!(5:8-5:13),
609                        },
610                        SimpleName {
611                            symbol: "a".to_string(),
612                            extent: optional_extent!(5:14-5:15),
613                        }
614                    ]
615                },
616                external_simple: DeclUse {
617                    declared_symbols: vec![],
618                    used_symbols: vec![],
619                },
620            }
621        );
622    }
623
624    #[test]
625    fn finds_names_in_decorated_definition() {
626        let code = r#"from metrics import measure_latency
627
628@measure_latency
629def get_items(key):
630    return client.post({ "key": key}).items
631"#;
632        let names = Python::find_names(code).expect("Cannot find names in test code");
633
634        assert_eq!(
635            names,
636            Names {
637                fully_qualified: DeclUse {
638                    declared_symbols: vec![FullyQualifiedName {
639                        source: sym!(metrics),
640                        symbol: sym!(measure_latency),
641                        extent: extent!(0:20-0:35),
642                    }],
643                    used_symbols: vec![FullyQualifiedName {
644                        source: sym!(metrics),
645                        symbol: sym!(measure_latency),
646                        extent: extent!(2:1-2:16),
647                    }]
648                },
649                simple: DeclUse {
650                    declared_symbols: vec![
651                        SimpleName {
652                            symbol: "metrics".to_string(),
653                            extent: optional_extent!(0:5-0:12),
654                        },
655                        SimpleName {
656                            symbol: "measure_latency".to_string(),
657                            extent: optional_extent!(0:20-0:35),
658                        },
659                        SimpleName {
660                            symbol: "get_items".to_string(),
661                            extent: optional_extent!(3:4-3:13),
662                        },
663                        SimpleName {
664                            symbol: "key".to_string(),
665                            extent: optional_extent!(3:14-3:17),
666                        }
667                    ],
668                    used_symbols: vec![
669                        SimpleName {
670                            symbol: "metrics".to_string(),
671                            extent: optional_extent!(),
672                        },
673                        SimpleName {
674                            symbol: "measure_latency".to_string(),
675                            extent: optional_extent!(2:1-2:16),
676                        },
677                        SimpleName {
678                            symbol: "client".to_string(),
679                            extent: optional_extent!(4:11-4:17),
680                        },
681                        SimpleName {
682                            symbol: "post".to_string(),
683                            extent: optional_extent!(4:18-4:22),
684                        },
685                        SimpleName {
686                            symbol: "key".to_string(),
687                            extent: optional_extent!(4:32-4:35),
688                        },
689                        SimpleName {
690                            symbol: "items".to_string(),
691                            extent: optional_extent!(4:38-4:43),
692                        },
693                    ]
694                },
695                external_simple: DeclUse {
696                    declared_symbols: vec![
697                        SimpleName {
698                            symbol: "metrics".to_string(),
699                            extent: optional_extent!(0:5-0:12),
700                        },
701                        SimpleName {
702                            symbol: "measure_latency".to_string(),
703                            extent: optional_extent!(0:20-0:35),
704                        }
705                    ],
706                    used_symbols: vec![
707                        SimpleName {
708                            symbol: "metrics".to_string(),
709                            extent: optional_extent!(),
710                        },
711                        SimpleName {
712                            symbol: "measure_latency".to_string(),
713                            extent: optional_extent!(2:1-2:16),
714                        },
715                        SimpleName {
716                            symbol: "client".to_string(),
717                            extent: optional_extent!(4:11-4:17),
718                        },
719                        SimpleName {
720                            symbol: "post".to_string(),
721                            extent: optional_extent!(4:18-4:22),
722                        },
723                        SimpleName {
724                            symbol: "items".to_string(),
725                            extent: optional_extent!(4:38-4:43),
726                        },
727                    ]
728                },
729            }
730        );
731    }
732
733    #[test]
734    fn finds_names_in_try_catch() {
735        let code = r#"try:
736  3 / 0
737except ZeroDivisionError as e:
738  print(e)
739finally:
740  a = 5
741
742print(a)
743"#;
744        let names = Python::find_names(code).expect("Cannot find names in test code");
745
746        assert_eq!(
747            names,
748            Names {
749                fully_qualified: DeclUse {
750                    declared_symbols: vec![],
751                    used_symbols: vec![]
752                },
753                simple: DeclUse {
754                    declared_symbols: vec![
755                        SimpleName {
756                            symbol: "e".to_string(),
757                            extent: optional_extent!(2:28-2:29),
758                        },
759                        SimpleName {
760                            symbol: "a".to_string(),
761                            extent: optional_extent!(5:2-5:3),
762                        }
763                    ],
764                    used_symbols: vec![
765                        SimpleName {
766                            symbol: "ZeroDivisionError".to_string(),
767                            extent: optional_extent!(2:7-2:24),
768                        },
769                        SimpleName {
770                            symbol: "print".to_string(),
771                            extent: optional_extent!(3:2-3:7),
772                        },
773                        SimpleName {
774                            symbol: "e".to_string(),
775                            extent: optional_extent!(3:8-3:9),
776                        },
777                        SimpleName {
778                            symbol: "print".to_string(),
779                            extent: optional_extent!(7:0-7:5),
780                        },
781                        SimpleName {
782                            symbol: "a".to_string(),
783                            extent: optional_extent!(7:6-7:7),
784                        }
785                    ]
786                },
787                external_simple: DeclUse {
788                    declared_symbols: vec![],
789                    used_symbols: vec![SimpleName {
790                        symbol: "ZeroDivisionError".to_string(),
791                        extent: optional_extent!(2:7-2:24),
792                    }]
793                },
794            }
795        );
796    }
797
798    #[test]
799    fn finds_names_in_generator_argument() {
800        let code =
801            r#"some_function(arg for arg in filter(is_valid, variants) if arg[0].startswith("a"))"#;
802
803        let names = Python::find_names(code).expect("Cannot find names in test code");
804
805        assert_eq!(
806            names,
807            Names {
808                fully_qualified: DeclUse {
809                    declared_symbols: vec![],
810                    used_symbols: vec![]
811                },
812                simple: DeclUse {
813                    declared_symbols: vec![SimpleName {
814                        symbol: "arg".to_string(),
815                        extent: optional_extent!(0:22-0:25),
816                    }],
817                    used_symbols: vec![
818                        SimpleName {
819                            symbol: "some_function".to_string(),
820                            extent: optional_extent!(0:0-0:13),
821                        },
822                        SimpleName {
823                            symbol: "filter".to_string(),
824                            extent: optional_extent!(0:29-0:35),
825                        },
826                        SimpleName {
827                            symbol: "is_valid".to_string(),
828                            extent: optional_extent!(0:36-0:44),
829                        },
830                        SimpleName {
831                            symbol: "variants".to_string(),
832                            extent: optional_extent!(0:46-0:54),
833                        },
834                        SimpleName {
835                            symbol: "arg".to_string(),
836                            extent: optional_extent!(0:59-0:62),
837                        },
838                        SimpleName {
839                            symbol: "startswith".to_string(),
840                            extent: optional_extent!(0:66-0:76),
841                        },
842                        SimpleName {
843                            symbol: "arg".to_string(),
844                            extent: optional_extent!(0:14-0:17),
845                        },
846                    ]
847                },
848                external_simple: DeclUse {
849                    declared_symbols: vec![],
850                    used_symbols: vec![
851                        SimpleName {
852                            symbol: "some_function".to_string(),
853                            extent: optional_extent!(0:0-0:13),
854                        },
855                        SimpleName {
856                            symbol: "is_valid".to_string(),
857                            extent: optional_extent!(0:36-0:44),
858                        },
859                        SimpleName {
860                            symbol: "variants".to_string(),
861                            extent: optional_extent!(0:46-0:54),
862                        },
863                        SimpleName {
864                            symbol: "startswith".to_string(),
865                            extent: optional_extent!(0:66-0:76),
866                        },
867                    ]
868                },
869            }
870        );
871    }
872
873    #[test]
874    fn finds_names_in_legacy_print_statement() {
875        let code = r#"print << "a = ", a"#;
876
877        let names = Python::find_names(code).expect("Cannot find names in test code");
878
879        assert_eq!(
880            names,
881            Names {
882                fully_qualified: DeclUse {
883                    declared_symbols: vec![],
884                    used_symbols: vec![]
885                },
886                simple: DeclUse {
887                    declared_symbols: vec![],
888                    used_symbols: vec![
889                        SimpleName {
890                            symbol: "print".to_string(),
891                            extent: optional_extent!(0:0-0:5),
892                        },
893                        SimpleName {
894                            symbol: "a".to_string(),
895                            extent: optional_extent!(0:17-0:18),
896                        }
897                    ]
898                },
899                external_simple: DeclUse {
900                    declared_symbols: vec![],
901                    used_symbols: vec![SimpleName {
902                        symbol: "a".to_string(),
903                        extent: optional_extent!(0:17-0:18),
904                    }]
905                },
906            }
907        );
908    }
909
910    #[test]
911    fn finds_boto3_resource() {
912        let code = r#"import boto3
913s3 = boto3.resource('s3')
914kinesis = boto3.client('kinesis', region='us-west-2')
915bucket = s3.Bucket('name')
916kinesis.put_record(
917    StreamName='some-stream',
918    Data=Bucket.name)
919"#;
920
921        let names = Python::find_names(code).expect("Cannot find names in test code");
922
923        assert_eq!(
924            names,
925            Names {
926                fully_qualified: DeclUse {
927                    declared_symbols: vec![FullyQualifiedName {
928                        source: sym!(boto3),
929                        symbol: vec![],
930                        extent: extent!(0:7-0:12)
931                    }],
932                    used_symbols: vec![
933                        FullyQualifiedName {
934                            source: sym!(boto3),
935                            symbol: vec![],
936                            extent: extent!(1:5-1:10)
937                        },
938                        FullyQualifiedName {
939                            source: sym!(boto3),
940                            symbol: sym!(resource),
941                            extent: extent!(1:11-1:19)
942                        },
943                        FullyQualifiedName {
944                            source: sym!(boto3),
945                            symbol: vec![],
946                            extent: extent!(2:10-2:15)
947                        },
948                        FullyQualifiedName {
949                            source: sym!(boto3),
950                            symbol: sym!(client),
951                            extent: extent!(2:16-2:22)
952                        },
953                        FullyQualifiedName {
954                            source: sym!(S3),
955                            symbol: sym!(ServiceResource.Bucket),
956                            extent: extent!(3:12-3:18)
957                        },
958                        FullyQualifiedName {
959                            source: sym!(Kinesis),
960                            symbol: sym!(Client.put_record),
961                            extent: extent!(4:8-4:18)
962                        }
963                    ]
964                },
965                simple: DeclUse {
966                    declared_symbols: vec![
967                        SimpleName {
968                            symbol: "boto3".to_string(),
969                            extent: optional_extent!(0:7-0:12)
970                        },
971                        SimpleName {
972                            symbol: "s3".to_string(),
973                            extent: optional_extent!(1:0-1:2)
974                        },
975                        SimpleName {
976                            symbol: "kinesis".to_string(),
977                            extent: optional_extent!(2:0-2:7)
978                        },
979                        SimpleName {
980                            symbol: "bucket".to_string(),
981                            extent: optional_extent!(3:0-3:6)
982                        }
983                    ],
984                    used_symbols: vec![
985                        SimpleName {
986                            symbol: "boto3".to_string(),
987                            extent: optional_extent!(1:5-1:10)
988                        },
989                        SimpleName {
990                            symbol: "boto3".to_string(),
991                            extent: None
992                        },
993                        SimpleName {
994                            symbol: "resource".to_string(),
995                            extent: optional_extent!(1:11-1:19)
996                        },
997                        SimpleName {
998                            symbol: "boto3".to_string(),
999                            extent: optional_extent!(2:10-2:15)
1000                        },
1001                        SimpleName {
1002                            symbol: "boto3".to_string(),
1003                            extent: None
1004                        },
1005                        SimpleName {
1006                            symbol: "client".to_string(),
1007                            extent: optional_extent!(2:16-2:22)
1008                        },
1009                        SimpleName {
1010                            symbol: "s3".to_string(),
1011                            extent: optional_extent!(3:9-3:11)
1012                        },
1013                        SimpleName {
1014                            symbol: "S3".to_string(),
1015                            extent: None
1016                        },
1017                        SimpleName {
1018                            symbol: "ServiceResource".to_string(),
1019                            extent: None
1020                        },
1021                        SimpleName {
1022                            symbol: "Bucket".to_string(),
1023                            extent: optional_extent!(3:12-3:18)
1024                        },
1025                        SimpleName {
1026                            symbol: "kinesis".to_string(),
1027                            extent: optional_extent!(4:0-4:7)
1028                        },
1029                        SimpleName {
1030                            symbol: "Kinesis".to_string(),
1031                            extent: None
1032                        },
1033                        SimpleName {
1034                            symbol: "Client".to_string(),
1035                            extent: None
1036                        },
1037                        SimpleName {
1038                            symbol: "put_record".to_string(),
1039                            extent: optional_extent!(4:8-4:18)
1040                        },
1041                        SimpleName {
1042                            symbol: "Bucket".to_string(),
1043                            extent: optional_extent!(6:9-6:15)
1044                        },
1045                        SimpleName {
1046                            symbol: "name".to_string(),
1047                            extent: optional_extent!(6:16-6:20)
1048                        }
1049                    ]
1050                },
1051                external_simple: DeclUse {
1052                    declared_symbols: vec![SimpleName {
1053                        symbol: "boto3".to_string(),
1054                        extent: optional_extent!(0:7-0:12)
1055                    }],
1056                    used_symbols: vec![
1057                        SimpleName {
1058                            symbol: "boto3".to_string(),
1059                            extent: optional_extent!(1:5-1:10)
1060                        },
1061                        SimpleName {
1062                            symbol: "boto3".to_string(),
1063                            extent: None
1064                        },
1065                        SimpleName {
1066                            symbol: "resource".to_string(),
1067                            extent: optional_extent!(1:11-1:19)
1068                        },
1069                        SimpleName {
1070                            symbol: "boto3".to_string(),
1071                            extent: optional_extent!(2:10-2:15)
1072                        },
1073                        SimpleName {
1074                            symbol: "boto3".to_string(),
1075                            extent: None
1076                        },
1077                        SimpleName {
1078                            symbol: "client".to_string(),
1079                            extent: optional_extent!(2:16-2:22)
1080                        },
1081                        SimpleName {
1082                            symbol: "S3".to_string(),
1083                            extent: None
1084                        },
1085                        SimpleName {
1086                            symbol: "ServiceResource".to_string(),
1087                            extent: None
1088                        },
1089                        SimpleName {
1090                            symbol: "Bucket".to_string(),
1091                            extent: optional_extent!(3:12-3:18)
1092                        },
1093                        SimpleName {
1094                            symbol: "Kinesis".to_string(),
1095                            extent: None
1096                        },
1097                        SimpleName {
1098                            symbol: "Client".to_string(),
1099                            extent: None
1100                        },
1101                        SimpleName {
1102                            symbol: "put_record".to_string(),
1103                            extent: optional_extent!(4:8-4:18)
1104                        },
1105                        SimpleName {
1106                            symbol: "Bucket".to_string(),
1107                            extent: optional_extent!(6:9-6:15)
1108                        },
1109                        SimpleName {
1110                            symbol: "name".to_string(),
1111                            extent: optional_extent!(6:16-6:20)
1112                        }
1113                    ]
1114                }
1115            }
1116        );
1117    }
1118
1119    #[test]
1120    fn finds_local_fields_through_self_parameter() {
1121        let code = r#"class Person:
1122  def __init__(self, name):
1123    self.name = name
1124
1125  def hello(self):
1126    return f"Hello, {self.name}"
1127"#;
1128
1129        let names = Python::find_names(code).expect("Cannot find names in test code");
1130
1131        assert_eq!(
1132            names,
1133            Names {
1134                fully_qualified: DeclUse {
1135                    declared_symbols: vec![],
1136                    used_symbols: vec![]
1137                },
1138                simple: DeclUse {
1139                    declared_symbols: vec![
1140                        SimpleName {
1141                            symbol: "Person".to_string(),
1142                            extent: optional_extent!(0:6-0:12),
1143                        },
1144                        SimpleName {
1145                            symbol: "__init__".to_string(),
1146                            extent: optional_extent!(1:6-1:14),
1147                        },
1148                        SimpleName {
1149                            symbol: "self".to_string(),
1150                            extent: optional_extent!(1:15-1:19),
1151                        },
1152                        SimpleName {
1153                            symbol: "name".to_string(),
1154                            extent: optional_extent!(1:21-1:25),
1155                        },
1156                        SimpleName {
1157                            symbol: "hello".to_string(),
1158                            extent: optional_extent!(4:6-4:11),
1159                        },
1160                        SimpleName {
1161                            symbol: "self".to_string(),
1162                            extent: optional_extent!(4:12-4:16),
1163                        }
1164                    ],
1165                    used_symbols: vec![
1166                        SimpleName {
1167                            symbol: "self".to_string(),
1168                            extent: optional_extent!(2:4-2:8),
1169                        },
1170                        SimpleName {
1171                            symbol: "Person".to_string(),
1172                            extent: None
1173                        },
1174                        SimpleName {
1175                            symbol: "name".to_string(),
1176                            extent: optional_extent!(2:9-2:13),
1177                        },
1178                        SimpleName {
1179                            symbol: "name".to_string(),
1180                            extent: optional_extent!(2:16-2:20),
1181                        }
1182                    ]
1183                },
1184                external_simple: DeclUse {
1185                    declared_symbols: vec![],
1186                    used_symbols: vec![]
1187                }
1188            }
1189        );
1190    }
1191
1192    #[test]
1193    fn finds_image_open() {
1194        let code = r#"from PIL import Image
1195
1196image = Image.open("C:\Users\System-Pc\Desktop\python.png")
1197"#;
1198
1199        let names = Python::find_names(code).expect("Cannot find names in test code");
1200
1201        assert_eq!(
1202            names,
1203            Names {
1204                fully_qualified: DeclUse {
1205                    declared_symbols: vec![FullyQualifiedName {
1206                        source: sym!(PIL),
1207                        symbol: sym!(Image),
1208                        extent: extent!(0:16-0:21),
1209                    }],
1210                    used_symbols: vec![
1211                        FullyQualifiedName {
1212                            source: sym!(PIL),
1213                            symbol: sym!(Image),
1214                            extent: extent!(2:8-2:13),
1215                        },
1216                        FullyQualifiedName {
1217                            source: sym!(PIL),
1218                            symbol: sym!(Image.open),
1219                            extent: extent!(2:14-2:18),
1220                        }
1221                    ]
1222                },
1223                simple: DeclUse {
1224                    declared_symbols: vec![
1225                        SimpleName {
1226                            symbol: "PIL".to_string(),
1227                            extent: optional_extent!(0:5-0:8),
1228                        },
1229                        SimpleName {
1230                            symbol: "Image".to_string(),
1231                            extent: optional_extent!(0:16-0:21),
1232                        },
1233                        SimpleName {
1234                            symbol: "image".to_string(),
1235                            extent: optional_extent!(2:0-2:5),
1236                        }
1237                    ],
1238                    used_symbols: vec![
1239                        SimpleName {
1240                            symbol: "PIL".to_string(),
1241                            extent: None
1242                        },
1243                        SimpleName {
1244                            symbol: "Image".to_string(),
1245                            extent: optional_extent!(2:8-2:13),
1246                        },
1247                        SimpleName {
1248                            symbol: "PIL".to_string(),
1249                            extent: None
1250                        },
1251                        SimpleName {
1252                            symbol: "Image".to_string(),
1253                            extent: None
1254                        },
1255                        SimpleName {
1256                            symbol: "open".to_string(),
1257                            extent: optional_extent!(2:14-2:18),
1258                        }
1259                    ]
1260                },
1261                external_simple: DeclUse {
1262                    declared_symbols: vec![
1263                        SimpleName {
1264                            symbol: "PIL".to_string(),
1265                            extent: optional_extent!(0:5-0:8),
1266                        },
1267                        SimpleName {
1268                            symbol: "Image".to_string(),
1269                            extent: optional_extent!(0:16-0:21),
1270                        }
1271                    ],
1272                    used_symbols: vec![
1273                        SimpleName {
1274                            symbol: "PIL".to_string(),
1275                            extent: None
1276                        },
1277                        SimpleName {
1278                            symbol: "Image".to_string(),
1279                            extent: optional_extent!(2:8-2:13),
1280                        },
1281                        SimpleName {
1282                            symbol: "PIL".to_string(),
1283                            extent: None
1284                        },
1285                        SimpleName {
1286                            symbol: "Image".to_string(),
1287                            extent: None
1288                        },
1289                        SimpleName {
1290                            symbol: "open".to_string(),
1291                            extent: optional_extent!(2:14-2:18),
1292                        }
1293                    ]
1294                }
1295            }
1296        );
1297    }
1298
1299    #[test]
1300    fn finds_relevant_snippet() {
1301        let code = r#"# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
1302# SPDX-License-Identifier: MIT-0
1303
1304from aws_cdk import (
1305    core as cdk,
1306    aws_cognito as cognito
1307)
1308from aws_cdk.aws_cognito import CfnUserPool, CfnUserPoolClient
1309
1310
1311class CognitoStack(cdk.Stack):
1312
1313    def __init__(self, scope: cdk.Construct, construct_id: str, admin_group_name="apiAdmins", **kwargs) -> None:
1314        super().__init__(scope, construct_id, **kwargs)
1315
1316        user_pool = CfnUserPool(
1317            self, "UserPoolInstance",
1318            user_pool_name=self.stack_name+"-UserPool",
1319            admin_create_user_config={'allowAdminCreateUserOnly': False},
1320            schema=[
1321                {
1322                    'attribute_data_type': 'String',
1323                    'name': 'name',
1324                    'mutable': True,
1325                    'required': True,
1326                },
1327                {
1328                    'attribute_data_type': 'String',
1329                    'name': 'email',
1330                    'mutable': True,
1331                    'required': True,
1332                },
1333            ],
1334            username_attributes=['email'],
1335        )
1336        user_pool_client = CfnUserPoolClient(
1337            self, "UserPoolClientInstance",
1338            client_name=self.stack_name+"UserPoolClient",
1339            user_pool_id=user_pool.ref,
1340            explicit_auth_flows=['ALLOW_USER_PASSWORD_AUTH','ALLOW_USER_SRP_AUTH','ALLOW_REFRESH_TOKEN_AUTH'],
1341            generate_secret=False,
1342            prevent_user_existence_errors='ENABLED',
1343            refresh_token_validity=30,
1344            supported_identity_providers=['COGNITO'],
1345            allowed_o_auth_flows_user_pool_client=True,
1346            allowed_o_auth_flows=['code'],
1347            allowed_o_auth_scopes=['email','openid'],
1348            callback_ur_ls=['http://localhost'],
1349        )
1350        user_pool_admin_group = cognito.CfnUserPoolGroup(
1351            self, "AdminUserPoolGroupInstance",
1352            user_pool_id=user_pool.ref,
1353            group_name=admin_group_name,
1354            description="User group for API Administrators",
1355            precedence=0
1356        )
1357        stack_name_prefix=self.stack_name.replace('-', '')
1358        cdk.CfnOutput(self, stack_name_prefix+'UserPool', export_name=stack_name_prefix+'UserPool', value=user_pool.ref, description='Cognito User Pool ID')
1359        cdk.CfnOutput(self, stack_name_prefix+'UserPoolClient', export_name=stack_name_prefix+'UserPoolClient', value=user_pool_client.ref, description='Cognito User Pool Application Client ID')
1360        cdk.CfnOutput(self, stack_name_prefix+'UserPoolAdminGroupName', export_name=stack_name_prefix+'UserPoolAdminGroupName', value=user_pool_admin_group.group_name, description='User Pool group name for API administrators')
1361        cdk.CfnOutput(self, stack_name_prefix+'CognitoLoginURL', export_name=stack_name_prefix+'CognitoLoginURL', value=f'https://${user_pool_client.ref}.auth.{cdk.Aws.REGION}.amazoncognito.com/login?client_id={user_pool_client.ref}&response_type=code&redirect_uri=http://localhost', description='Cognito User Pool Application Client Hosted Login UI URL')
1362        cdk.CfnOutput(self, stack_name_prefix+'CognitoAuthCommand', export_name=stack_name_prefix+'CognitoAuthCommand', value=f'aws cognito-idp initiate-auth --auth-flow USER_PASSWORD_AUTH --client-id {user_pool_client.ref} --auth-parameters USERNAME=<username>,PASSWORD=<password>', description='AWS CLI command for Amazon Cognito User Pool authentication')"#;
1363
1364        let snippet = Python::find_relevant_scope(
1365            code,
1366            vec![RelevantName::FullyQualified {
1367                source: sym!(aws_cdk.aws_cognito),
1368                symbol: sym!(CfnUserPool),
1369            }],
1370        )
1371        .expect("Cannot find snippet in test code")
1372        .snippet;
1373
1374        assert_eq!(
1375            snippet,
1376            r#"def __init__(self, scope: cdk.Construct, construct_id: str, admin_group_name="apiAdmins", **kwargs) -> None:
1377        super().__init__(scope, construct_id, **kwargs)
1378
1379        user_pool = CfnUserPool(
1380            self, "UserPoolInstance",
1381            user_pool_name=self.stack_name+"-UserPool",
1382            admin_create_user_config={'allowAdminCreateUserOnly': False},
1383            schema=[
1384                {
1385                    'attribute_data_type': 'String',
1386                    'name': 'name',
1387                    'mutable': True,
1388                    'required': True,
1389                },
1390                {
1391                    'attribute_data_type': 'String',
1392                    'name': 'email',
1393                    'mutable': True,
1394                    'required': True,
1395                },
1396            ],
1397            username_attributes=['email'],
1398        )
1399        user_pool_client = CfnUserPoolClient(
1400            self, "UserPoolClientInstance",
1401            client_name=self.stack_name+"UserPoolClient",
1402            user_pool_id=user_pool.ref,
1403            explicit_auth_flows=['ALLOW_USER_PASSWORD_AUTH','ALLOW_USER_SRP_AUTH','ALLOW_REFRESH_TOKEN_AUTH'],
1404            generate_secret=False,
1405            prevent_user_existence_errors='ENABLED',
1406            refresh_token_validity=30,
1407            supported_identity_providers=['COGNITO'],
1408            allowed_o_auth_flows_user_pool_client=True,
1409            allowed_o_auth_flows=['code'],
1410            allowed_o_auth_scopes=['email','openid'],
1411            callback_ur_ls=['http://localhost'],
1412        )
1413        user_pool_admin_group = cognito.CfnUserPoolGroup(
1414            self, "AdminUserPoolGroupInstance",
1415            user_pool_id=user_pool.ref,
1416            group_name=admin_group_name,
1417            description="User group for API Administrators",
1418            precedence=0
1419        )
1420        stack_name_prefix=self.stack_name.replace('-', '')
1421        cdk.CfnOutput(self, stack_name_prefix+'UserPool', export_name=stack_name_prefix+'UserPool', value=user_pool.ref, description='Cognito User Pool ID')
1422        cdk.CfnOutput(self, stack_name_prefix+'UserPoolClient', export_name=stack_name_prefix+'UserPoolClient', value=user_pool_client.ref, description='Cognito User Pool Application Client ID')
1423        cdk.CfnOutput(self, stack_name_prefix+'UserPoolAdminGroupName', export_name=stack_name_prefix+'UserPoolAdminGroupName', value=user_pool_admin_group.group_name, description='User Pool group name for API administrators')
1424        cdk.CfnOutput(self, stack_name_prefix+'CognitoLoginURL', export_name=stack_name_prefix+'CognitoLoginURL', value=f'https://${user_pool_client.ref}.auth.{cdk.Aws.REGION}.amazoncognito.com/login?client_id={user_pool_client.ref}&response_type=code&redirect_uri=http://localhost', description='Cognito User Pool Application Client Hosted Login UI URL')
1425        cdk.CfnOutput(self, stack_name_prefix+'CognitoAuthCommand', export_name=stack_name_prefix+'CognitoAuthCommand', value=f'aws cognito-idp initiate-auth --auth-flow USER_PASSWORD_AUTH --client-id {user_pool_client.ref} --auth-parameters USERNAME=<username>,PASSWORD=<password>', description='AWS CLI command for Amazon Cognito User Pool authentication')"#
1426        );
1427    }
1428
1429    #[test]
1430    fn finds_names_in_type_annotation() {
1431        // Given
1432        let code = "a: T: Iterable[int] = []";
1433
1434        // When
1435        let names = Python::find_names(code).expect("Cannot find names in test code");
1436
1437        // Then
1438        assert_eq!(
1439            names,
1440            Names {
1441                fully_qualified: DeclUse {
1442                    declared_symbols: vec![],
1443                    used_symbols: vec![]
1444                },
1445                simple: DeclUse {
1446                    declared_symbols: vec![SimpleName {
1447                        symbol: "a".to_string(),
1448                        extent: optional_extent!(0:0-0:1)
1449                    }],
1450                    used_symbols: vec![
1451                        SimpleName {
1452                            symbol: "T".to_string(),
1453                            extent: optional_extent!(0:3-0:4)
1454                        },
1455                        SimpleName {
1456                            symbol: "Iterable".to_string(),
1457                            extent: optional_extent!(0:6-0:14)
1458                        },
1459                        SimpleName {
1460                            symbol: "int".to_string(),
1461                            extent: optional_extent!(0:15-0:18)
1462                        }
1463                    ]
1464                },
1465                external_simple: DeclUse {
1466                    declared_symbols: vec![],
1467                    used_symbols: vec![
1468                        SimpleName {
1469                            symbol: "T".to_string(),
1470                            extent: optional_extent!(0:3-0:4)
1471                        },
1472                        SimpleName {
1473                            symbol: "Iterable".to_string(),
1474                            extent: optional_extent!(0:6-0:14)
1475                        }
1476                    ]
1477                }
1478            }
1479        )
1480    }
1481}