# coding=utf-8

"""
A Diamond collector that collects memory usage of each process defined in it's
config file by matching them with their executable filepath or the process name.

Example config file ProcessMemoryCollector.conf

```
enabled=True
unit=kB
[process]
[[postgres]]
exe=^\/usr\/lib\/postgresql\/+d.+d\/bin\/postgres$
name=^postgres,^pg
```

exe and name are both lists of comma-separated regexps.
"""
import re

import diamond.collector, diamond.convertor
import psutil

class ProcessMemoryCollector(diamond.collector.Collector):

    def get_default_config_help(self):
        config_help = super(ProcessMemoryCollector,
                            self).get_default_config_help()
        config_help.update({
            'unit': 'The unit in which memory data is collected.',
            'process': ("A subcategory of settings inside of which each "
                        "collected process has it's configuration")
        })
        return config_help

    def get_default_config(self):
        """
        Default settings are:
            path: 'memory.process'
            unit: 'B'
        """
        config = super(ProcessMemoryCollector, self).get_default_config()
        config.update({
            'path': 'memory.process',
            'unit': 'B',
            'process': '',
            })
        return config

    def setup_config(self):
        """
        prepare self.processes, which is a descriptor dictionary in
        processgroup --> {
            exe: [regex],
            name: [regex],
            procs: [psutil.Process]
            }
        """
        self.processes = {}
        for process, cfg in self.config['process'].items():
            # first we build a dictionary with the process aliases and the
            #  matching regexps
            exe = cfg.get('exe', [])
            if not isinstance(exe, list):
                exe = [exe]
            exe = [re.compile(e) for e in exe]
            name = cfg.get('name', [])
            if not isinstance(name, list):
                name = [name]
            name = [re.compile(n) for n in name]
            self.processes[process] = {
                'exe': exe,
                'name': name,
                'procs': []
            }

    def filter_processes(self):
        """
        Populates self.processes[processname]['procs'] with the corresponding
        list of psutil.Process instances
        """
        def process_filter(proc, cfg):
            """
            Decides whether a process matches with a given process descriptor

            :param proc: a psutil.Process instance
            :param cfg: the dictionary from processes that describes with the
                process group we're testing for
            :return: True if it matches
            :rtype: bool
            """
            for exe in cfg['exe']:
                try:
                    if exe.search(proc.exe):
                        return True
                except psutil.AccessDenied:
                    break
            for name in cfg['name']:
                if name.search(proc.name):
                    return True
            return False

        for proc in psutil.process_iter():
            # filter and divide the system processes amongst the different
            #  process groups defined in the config file
            for procname, cfg in self.processes.items():
                if process_filter(proc, cfg):
                    cfg['procs'].append(proc)
                    break


    def collect(self):
        """
        Collects the RSS memory usage of each process defined under the
        `process` subsection of the config file
        """
        self.setup_config()
        self.filter_processes()
        unit = self.config['unit']
        for process, cfg in self.processes.items():
            # finally publish the results for each process group
            metric_name = "{}.rss".format(process)
            metric_value = diamond.convertor.binary.convert(
                sum(p.get_memory_info().rss for p in cfg['procs']),
                oldUnit='byte', newUnit=unit)
            # Publish Metric
            self.publish(metric_name, metric_value)

            metric_name = "{}.vms".format(process)
            metric_value = diamond.convertor.binary.convert(
                sum(p.get_memory_info().vms for p in cfg['procs']),
                oldUnit='byte', newUnit=unit)
            # Publish Metric
            self.publish(metric_name, metric_value)

@@ -1796,32 +1796,28 @@
     
-exe = cfg.get('exe', %5B%5D)
+proc = %7B'procs': %5B%5D%7D
 %0A   
@@ -1829,35 +1829,44 @@
     
-if not isinstance(
+for key in ('
 exe
+'
 , 
-list
+'name', 'cmdline'
 ):%0A 
@@ -1884,83 +1884,17 @@
     
-exe = %5Bexe%5D%0A            exe = %5Bre.compile(e) for e in exe%5D%0A            name
+proc%5Bkey%5D
  = c
@@ -1904,20 +1904,21 @@
 get(
-'name'
+key
 , %5B%5D)%0A
+    
     
@@ -1943,20 +1943,25 @@
 nstance(
-name
+proc%5Bkey%5D
 , list):
@@ -1981,20 +1981,34 @@
     
-name = %5Bname
+    proc%5Bkey%5D = %5Bproc%5Bkey%5D
 %5D%0A  
@@ -2017,20 +2017,29 @@
         
-name
+    proc%5Bkey%5D
  = %5Bre.c
@@ -2049,24 +2049,37 @@
 ile(
-n
+proc%5Bkey%5D
 ) for 
-n
+e
  in 
-name
+proc%5Bkey%5D
 %5D%0A  
@@ -2118,109 +2118,12 @@
 %5D = 
-%7B%0A                'exe': exe,%0A                'name': name,%0A                'procs': %5B%5D%0A            %7D
+proc
 %0A%0A  
@@ -2998,32 +2998,166 @@
     return True%0A
+            for cmdline in cfg%5B'cmdline'%5D:%0A                if cmdline.search(' '.join(proc.cmdline)):%0A                    return True%0A
             retu

check if cmdline of process match