How do I get all the tokens for a statement? #4

Activity

1. lexanth commented on Mar 23, 2023

   

   lexanth

   on Mar 23, 2023

   Owner

   More actions

   I've not used the getTokens method really, but it looks like it doesn't do deep searching, so may not be what you're after. The constants for token types are on Python3Parser though - e.g. Python3Parser.STRING.

   To actually investigate that sort of piece of code, I'd look at using the walk or visitor APIs to pick the type of element you're looking for first - e.g.

   const code = `from datetime import datetime, date
   `;
   const ast = parse(code);
   
   const visitor = createVisitor({
     visitImport_from(importNode) {
       const importSourceNode = importNode.dotted_name();
       const importedItems = importNode.import_as_names()?.import_as_name();
   
       console.log(importSourceNode?.text);
       console.log(importedItems?.map((item) => item.text));
     },
   });
   visitor.visit(ast);
   
   const imports = [];
   walk(
     {
       enterImport_from(importNode) {
         const importSourceNode = importNode.dotted_name();
         const importedItems = importNode.import_as_names()?.import_as_name();
   
         imports.push({
           source: importSourceNode?.text,
           items: importedItems.map((item) => item.text),
         });
       },
     },
     ast
   );
   console.log(imports);
   

   This snippet still doesn't handle the full complexity of python imports, so e.g. renamed imports would need something different, but this might set you on the right path, depending on what you're trying to do.
2. hirehamir commented on Mar 23, 2023

   

   hirehamir

   on Mar 23, 2023

   Author

   More actions

   That snippet definitely helps. Thanks for posting it.

   Do you by any chance have any suggestions on how to go about extracting a renamed import from import datetime as dt, for example?

   This is what I tried.

   import { createVisitor, parse, walk } from 'python-ast';
   
   interface Import {
     items: string[];
     source: string;
   }
   
   const analyze = (code: string) => {
     const ast = parse(code);
     const visitor = createVisitor({
       visitImport_name(importNode) {
         const importSourceNode = importNode.dotted_as_names();
   
         console.log('visitImport_name', importSourceNode?.text);
       },
       visitImport_stmt(importNode) {
         const importSourceNode = importNode.import_name();
         importSourceNode
           ?.dotted_as_names()
           .children?.forEach((child) => console.log('child', child.text));
         importSourceNode
           ?.dotted_as_names()
           .children?.forEach((child) => console.log('child', child.text));
   
         console.log('visitImport_stmt', importSourceNode?.text);
         console.log(
           'visitImport_stmt',
           importSourceNode
             ?.dotted_as_names()
             .dotted_as_name()
             .map((i) => i.text)
         );
       }
     });
     visitor.visit(ast);
   
     const imports: Import[] = [];
     walk(
       {
         enterImport_from(importNode) {
           const importSourceNode = importNode.dotted_name();
           const importedItems = importNode.import_as_names()?.import_as_name();
   
           imports.push({
             source: importSourceNode?.text || '',
             items: importedItems?.map((item) => item.text) || []
           });
         }
       },
       ast
     );
     return imports;
   };
   
   analyze('import datetime as dt\n');

   This was the output.

   child datetimeasdt
   child datetimeasdt
   visitImport_stmt importdatetimeasdt
   visitImport_stmt [ 'datetimeasdt' ]